Create previewer/modules.py

previewer/modules.py

@@ -0,0 +1,36 @@
+from torch import nn
+
+# Effnet 16x16 to 64x64 previewer
+class Previewer(nn.Module):
+    def __init__(self, c_in=16, c_hidden=512, c_out=3):
+        super().__init__()
+        self.blocks = nn.Sequential(
+            nn.Conv2d(c_in, c_hidden, kernel_size=1), # 36 channels to 512 channels
+            nn.GELU(),
+            nn.BatchNorm2d(c_hidden),
+
+            nn.Conv2d(c_hidden, c_hidden, kernel_size=3, padding=1),
+            nn.GELU(),
+            nn.BatchNorm2d(c_hidden),
+
+            nn.ConvTranspose2d(c_hidden, c_hidden//2, kernel_size=2, stride=2), # 16 -> 32
+            nn.GELU(),
+            nn.BatchNorm2d(c_hidden//2),
+
+            nn.Conv2d(c_hidden//2, c_hidden//2, kernel_size=3, padding=1),
+            nn.GELU(),
+            nn.BatchNorm2d(c_hidden//2),
+
+            nn.ConvTranspose2d(c_hidden//2, c_hidden//4, kernel_size=2, stride=2), # 32 -> 64
+            nn.GELU(),
+            nn.BatchNorm2d(c_hidden//4),
+
+            nn.Conv2d(c_hidden//4, c_hidden//4, kernel_size=3, padding=1),
+            nn.GELU(),
+            nn.BatchNorm2d(c_hidden//4),
+
+            nn.Conv2d(c_hidden//4, c_out, kernel_size=1),
+        )
+
+    def forward(self, x):
+        return self.blocks(x)