Create decoder.py

de_en/decoder.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class VideoDecoder(nn.Module):
+    def __init__(self, embed_dim=512, hidden_dims=[512, 256, 128, 64], out_channels=3):
+        super().__init__()
+        
+        self.fc = nn.Linear(embed_dim, hidden_dims[0] * 8 * 8)
+        
+        modules = []
+        for i in range(len(hidden_dims)-1):
+            modules.append(
+                nn.Sequential(
+                    nn.ConvTranspose2d(hidden_dims[i], hidden_dims[i+1],
+                                     kernel_size=3, stride=2, padding=1, output_padding=1),
+                    nn.BatchNorm2d(hidden_dims[i+1]),
+                    nn.LeakyReLU()
+                )
+            )
+        
+        self.decoder = nn.Sequential(*modules)
+        
+        self.final_layer = nn.Sequential(
+            nn.ConvTranspose2d(hidden_dims[-1], hidden_dims[-1],
+                              kernel_size=3, stride=2, padding=1, output_padding=1),
+            nn.BatchNorm2d(hidden_dims[-1]),
+            nn.LeakyReLU(),
+            nn.Conv2d(hidden_dims[-1], out_channels, kernel_size=3, padding=1),
+            nn.Tanh()
+        )
+        
+    def forward(self, z):
+        # z shape: (seq_len, batch, embed_dim) or (batch, embed_dim)
+        if z.dim() == 3:
+            batch_size, seq_len = z.size(1), z.size(0)
+            z = z.view(-1, z.size(2))  # flatten for linear layer
+        else:
+            batch_size, seq_len = z.size(0), 1
+        
+        x = self.fc(z)
+        x = x.view(-1, self.decoder[0][0].in_channels, 8, 8)
+        x = self.decoder(x)
+        x = self.final_layer(x)
+        
+        if seq_len > 1:
+            x = x.view(seq_len, batch_size, *x.shape[1:])
+        
+        return x