view training trials

models/view_encoder.py

@@ -22,42 +22,72 @@ class ViewEncoder(nn.Module):
         super(ViewEncoder, self).__init__()
         self.bottleneck_size = bottleneck_size
         
-        # Encoder
-        self.encoder = nn.Sequential(
-            nn.Linear(3*3, 64),
-            nn.ReLU(),
-            nn.Linear(64, 128),
-            nn.ReLU(),
-            nn.Linear(128, 256),
-            nn.ReLU(),
-            nn.Linear(256, 512),
-            nn.ReLU(),
-            nn.Linear(512, bottleneck_size),
-            nn.ReLU()
-        )
+        self.encoder = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(3*3, 64),
+                nn.BatchNorm1d(64),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(64, 128),
+                nn.BatchNorm1d(128),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(128, 256),
+                nn.BatchNorm1d(256),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(256, 512),
+                nn.BatchNorm1d(512),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(512, bottleneck_size),
+                nn.BatchNorm1d(bottleneck_size),)
+        ])
 
         # Decoder
-        self.decoder = nn.Sequential(
-            nn.Linear(bottleneck_size, 512),
-            nn.ReLU(),
-            nn.Linear(512, 256),
-            nn.ReLU(),
-            nn.Linear(256, 128),
-            nn.ReLU(),
-            nn.Linear(128, 64),
-            nn.ReLU(),
-            nn.Linear(64, 3*3),
-            nn.Tanh(),
-        )
+        self.decoder = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(bottleneck_size, 512),
+                nn.BatchNorm1d(512),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(512*2, 256),
+                nn.BatchNorm1d(256),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(256*2, 128),
+                nn.BatchNorm1d(128),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(128*2, 64),
+                nn.BatchNorm1d(64),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(64*2, 3*3),
+                nn.Tanh(),),
+        ])
     
     def forward(self, x):
         # Encode
-        matrix = self.encoder(x.view(-1, 3*3))
+        f1 = self.encoder[0](x.view(-1, 3*3))
+        f2 = self.encoder[1](f1)
+        f3 = self.encoder[2](f2)
+        f4 = self.encoder[3](f3)
+        matrix = self.encoder[4](f4)
+
+        d1 = self.decoder[0](matrix)
+        d2 = self.decoder[1](torch.cat((d1, f4), dim=1))
+        d3 = self.decoder[2](torch.cat((d2, f3), dim=1))
+        d4 = self.decoder[3](torch.cat((d3, f2), dim=1))
+        x = self.decoder[4](torch.cat((d4, f1), dim=1))
         
-        # Decode
-        x = self.decoder(matrix)
         x = x * np.pi
         return x.view(-1, 3, 3), matrix
+        
+        # # Decode
+        # x = self.decoder(matrix)
+        # x = x * np.pi
+        # return x.view(-1, 3, 3), matrix
     
 
 if __name__ == '__main__':