wip

shape_complexity/shape_complexity.py

@@ -142,7 +142,7 @@ class CONVVAE(nn.Module):
         super(CONVVAE, self).__init__()
 
         self.bottleneck = bottleneck
-        self.feature_dim = 128
+        self.feature_dim = 32
 
         self.conv1 = nn.Sequential(
             nn.Conv2d(1, 16, 5),
@@ -160,44 +160,65 @@ class CONVVAE(nn.Module):
             nn.MaxPool2d((2, 2), return_indices=True),  # -> 6x6x64
         )
         self.conv4 = nn.Sequential(
-            nn.Conv2d(64, 128, 5),
+            nn.Conv2d(64, 2 * self.bottleneck, 5),
             nn.ReLU(),
-            nn.MaxPool2d((2, 2), return_indices=True),  # -> 1x1x128
+            nn.MaxPool2d((2, 2), return_indices=True),  # -> 1x1x2*bottleneck
         )
 
         self.encode_mu = nn.Sequential(
-            nn.Flatten(), nn.Linear(self.feature_dim, self.bottleneck)
+            nn.Flatten(),
+            nn.Linear(
+                2 * self.bottleneck, self.bottleneck
+            ),  # TODO: maybe only FC from bn x bn
         )
         self.encode_logvar = nn.Sequential(
-            nn.Flatten(), nn.Linear(self.feature_dim, self.bottleneck)
+            nn.Flatten(), nn.Linear(2 * self.bottleneck, self.bottleneck)
         )
 
-        self.decode = nn.Sequential(
-            nn.MaxUnpool2d((2, 2)),
-            nn.ConvTranspose2d(128, 64, 5),
+        self.decode_linear = nn.Linear(self.bottleneck, 2 * self.bottleneck)
+
+        self.decode4 = nn.Sequential(
+            nn.ConvTranspose2d(2 * self.bottleneck, 64, 5),
             nn.ReLU(),
-            nn.MaxUnpool2d((2, 2)),
+        )
+        self.decode3 = nn.Sequential(
             nn.ConvTranspose2d(64, 32, 3),
             nn.ReLU(),
-            nn.MaxUnpool2d((2, 2)),
+        )
+        self.decode2 = nn.Sequential(
             nn.ConvTranspose2d(32, 16, 3),
             nn.ReLU(),
-            nn.MaxUnpool2d((2, 2)),
+        )
+        self.decode1 = nn.Sequential(
             nn.ConvTranspose2d(16, 1, 5),
             nn.Sigmoid(),
         )
 
     def encode(self, x):
-        x, indices = self.conv(x)
+        x, idx1 = self.conv1(x)
+        x, idx2 = self.conv2(x)
+        x, idx3 = self.conv3(x)
+        x, idx4 = self.conv4(x)
         mu = self.encode_mu(x)
         logvar = self.encode_logvar(x)
 
-        return mu, logvar
-
-    # def decode(self, z):
-    #     z = self.decode_linear(z)
-    #     # z = z.view(-1, 128, 1, 1)
-    #     return self.decode_conv(z)
+        return mu, logvar, (idx1, idx2, idx3, idx4)
+
+    def decode(self, z: Tensor, indexes: tuple):
+        (idx1, idx2, idx3, idx4) = indexes
+        z = self.decode_linear(z)
+        z = z.view((-1, 2 * self.bottleneck, 1, 1))
+        z = F.max_unpool2d(z, idx4, (2, 2))
+        z = self.decode4(z)
+        z = F.max_unpool2d(z, idx3, (2, 2))
+        z = self.decode3(z)
+        z = F.max_unpool2d(z, idx2, (2, 2))
+        z = self.decode2(z)
+        z = F.max_unpool2d(z, idx1, (2, 2))
+        z = self.decode1(z)
+        # z = z.view(-1, 128, 1, 1)
+        # return self.decode_conv(z)
+        return z
 
     def reparameterize(self, mu, logvar):
         std = torch.exp(0.5 * logvar)
@@ -205,9 +226,9 @@ class CONVVAE(nn.Module):
         return mu + eps * std
 
     def forward(self, x):
-        mu, logvar = self.encode(x)
+        mu, logvar, indexes = self.encode(x)
         z = self.reparameterize(mu, logvar)
-        return self.decode(z), mu, logvar
+        return self.decode(z, indexes), mu, logvar
 
     def loss(self, recon_x, x, mu, logvar):
         """https://github.com/pytorch/examples/blob/main/vae/main.py"""