direct training

models/direct.py

@@ -0,0 +1,29 @@
+import torch
+from torch import nn
+from models.convnext_orig import ConvNeXt, CNBlock, CNBlockConfig
+from typing import Any, Callable, List, Optional, Sequence
+
+
+class DirectModel(nn.Module):
+    
+    def __init__(self, block_setting: List[CNBlockConfig], kernel_size: int = 7, input_features: int = 16, device: str = 'cuda'):
+        super(DirectModel, self).__init__()
+
+        self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, input_features=input_features, device=device)
+        self.decoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, device=device, reversed=True)
+
+        layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
+        self.mean_branch = nn.Sequential(*layers).to(device)
+        layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
+        self.var_branch = nn.Sequential(*layers).to(device)
+
+        self.flatten_mean = nn.Flatten()
+        self.flatten_var = nn.Flatten()
+
+
+    def forward(self, inputs):
+        body, template = inputs
+
+        x = self.encoder(x)
+
+        return mean