Refine minist dump.

tools/show.py

@@ -8,7 +8,7 @@ import numpy as np
 import os
 
 
-def DumpTensorToImage(tensor, name, forceSquare=True, scale=1.0, Contrast=None, GridValue=None):
+def DumpTensorToImage(tensor, name, forceSquare=False, scale=1.0, Contrast=None, GridValue=None):
     if len(tensor.shape) != 2 and len(tensor.shape) != 1 and len(tensor.shape) != 3:
         raise ("Error input dims")
     if ("." not in name) or (name.split(".")[-1] not in {"jpg", "png", "bmp"}):

unsuper/minist.py

@@ -34,15 +34,15 @@ class ConvNet(nn.Module):
         super(ConvNet, self).__init__()
         self.conv1 = nn.Conv2d(1, 8, 5, 1, 0)
         self.pool = nn.MaxPool2d(2, 2)
-        self.conv2 = nn.Conv2d(8, 8, 5, 1, 0)
-        self.fc1 = nn.Linear(8 * 4 * 4, 10)
+        self.conv2 = nn.Conv2d(8, 1, 5, 1, 0)
+        self.fc1 = nn.Linear(1 * 4 * 4, 10)
         # self.fc2 = nn.Linear(120, 84)
         # self.fc3 = nn.Linear(84, 10)
 
     def forward(self, x):
         x = self.pool(F.relu(self.conv1(x)))
         x = self.pool(F.relu(self.conv2(x)))
-        x = x.view(-1, 8 * 4 * 4)
+        x = x.view(x.shape[0], -1)
         # x = F.relu(self.fc1(x))
         # x = F.relu(self.fc2(x))
         # x = self.fc3(x)
@@ -64,15 +64,17 @@ class ConvNet(nn.Module):
         x = self.pool(F.relu(x))
         x = self.conv2(x)
         w = self.conv2.weight
-        show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv2_weight.png")
+        show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv2_weight.png", Contrast=[-1.0, 1.0])
 
-        show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]).cpu(), "conv2_output.png")
+        show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]).cpu(), "conv2_output.png", Contrast=[-1.0, 1.0])
         x = self.pool(F.relu(x))
-        x = x.view(-1, 8 * 4 * 4)
+        show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]).cpu(), "pool_output.png", Contrast=[-1.0, 1.0])
+        pool_shape = x.shape
+        x = x.view(x.shape[0], -1)
         x = self.fc1(x)
-
-        show.DumpTensorToImage(self.fc1.weight.view(-1, 16, 8).permute(2, 0, 1), "fc_weight.png")
-
+        show.DumpTensorToImage(
+            self.fc1.weight.view(-1, pool_shape[2], pool_shape[3]), "fc_weight.png", Contrast=[-1.0, 1.0]
+        )
         show.DumpTensorToImage(x.view(-1).cpu(), "fc_output.png")
 
         criterion = nn.CrossEntropyLoss()
@@ -84,7 +86,7 @@ class ConvNet(nn.Module):
         show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv1_weight_grad.png")
         w = self.conv2.weight.grad
         show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), "conv2_weight_grad.png")
-        show.DumpTensorToImage(self.fc1.weight.grad.view(-1, 16, 8).permute(2, 0, 1), "fc_weight_grad.png")
+        show.DumpTensorToImage(self.fc1.weight.grad.view(-1, pool_shape[2], pool_shape[3]), "fc_weight_grad.png")
 
 
 model = ConvNet().to(device)