Refine to dump heat image.

tools/show.py

@@ -8,7 +8,7 @@ import numpy as np
 import os
 
 
-def DumpTensorToImage(tensor, name, forceSquare=True, scale=1.0, AutoContrast=True, GridValue=0):
+def DumpTensorToImage(tensor, name, forceSquare=True, 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"}):
@@ -17,27 +17,36 @@ def DumpTensorToImage(tensor, name, forceSquare=True, scale=1.0, AutoContrast=Tr
     if len(tensor.shape) == 3:
         channel = tensor.shape[0]
         x = math.ceil((channel) ** 0.5)
-        tensor = F.pad(tensor, (0, 0, 0, 0, 0, x * x - channel), mode="constant", value=0)
+        calc = tensor.reshape((channel, tensor.shape[1] * tensor.shape[2]))
-        calc = tensor.reshape((x * x, tensor.shape[1] * tensor.shape[2]))
+        if not Contrast:
-        if AutoContrast:
             tensormax = calc.max(1)[0]
             tensormin = calc.min(1)[0]
+        else:
+            tensormax = Contrast[1]
+            tensormin = Contrast[0]
         calc = calc.transpose(1, 0)
-            calc = ((calc - tensormin) / (tensormax - tensormin)) * 255
+        calc = ((calc - tensormin) / (tensormax - tensormin)) * 255.0
         calc = calc.transpose(1, 0)
+        calc = calc.reshape((channel, tensor.shape[1], tensor.shape[2]))
+        if not GridValue:
+            GridValue = 128.0
+        calc = F.pad(calc, (0, 0, 0, 0, 0, x * x - channel), mode="constant", value=GridValue)
         calc = calc.reshape((x, x, tensor.shape[1], tensor.shape[2]))
         calc = F.pad(calc, (0, 1, 0, 1, 0, 0), mode="constant", value=GridValue)
         tensor = calc.permute((0, 2, 1, 3))
         tensor = tensor.reshape((x * tensor.shape[1], x * tensor.shape[3]))
-        DumpTensorToImage(tensor, name, forceSquare=False, scale=scale, AutoContrast=False)
+        DumpTensorToImage(tensor, name, forceSquare=False, scale=scale, Contrast=[0.0, 255.0], GridValue=GridValue)
         return
 
     tensor = tensor.float()
-    if AutoContrast:
+    if not Contrast:
         maxv = torch.max(tensor)
         minv = torch.min(tensor)
-        tensor = ((tensor - minv) / (maxv - minv)) * 255
+    else:
-    img = tensor.byte().cpu().numpy()
+        maxv = Contrast[1]
+        minv = Contrast[0]
+    tensor = ((tensor - minv) / (maxv - minv)) * 255.0
+    img = tensor.detach().cpu().numpy()
     srp = img.shape
 
     if len(srp) == 1:  # 1D的数据自动折叠成2D图像
@@ -51,13 +60,21 @@ def DumpTensorToImage(tensor, name, forceSquare=True, scale=1.0, AutoContrast=Tr
     if scale != 1.0:
         img = cv2.resize(img, [int(srp[0] * scale), int(srp[1] * scale)])
         srp = img.shape
-    cv2.imwrite(name, img)
+
+    img = img * (-1)
+    img = img + 255
+    img[img < 0] = 0
+    img = np.nan_to_num(img, nan=0.0)
+    img[img > 255] = 255
+    imgs = img.astype(np.uint8)
+    imgs = cv2.applyColorMap(imgs, cv2.COLORMAP_JET)
+    cv2.imwrite(name, imgs)
 
 
 def DumpTensorToLog(tensor, name="log"):
     shape = tensor.shape
     f = open(name, "w")
-    data = tensor.reshape([-1]).float().cpu().numpy().tolist()
+    data = tensor.reshape([-1]).float().cpu().detach().numpy().tolist()
     for d in data:
         f.writelines("%s" % d + os.linesep)
     f.close()

unsuper/minist.py

@@ -14,11 +14,11 @@ torch.manual_seed(seed)
 torch.cuda.manual_seed_all(seed)
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-device = torch.device("mps")
+# device = torch.device("mps")
 
 # Hyper-parameters
 num_epochs = 1
-batch_size = 1
+batch_size = 256
 learning_rate = 0.001
 
 transform = transforms.Compose([transforms.ToTensor()])
@@ -32,17 +32,17 @@ test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, s
 class ConvNet(nn.Module):
     def __init__(self):
         super(ConvNet, self).__init__()
-        self.conv1 = nn.Conv2d(1, 8, 3, 1, 1)
+        self.conv1 = nn.Conv2d(1, 8, 5, 1, 0)
         self.pool = nn.MaxPool2d(2, 2)
-        self.conv2 = nn.Conv2d(8, 8, 5)
+        self.conv2 = nn.Conv2d(8, 8, 5, 1, 0)
-        self.fc1 = nn.Linear(8 * 5 * 5, 10)
+        self.fc1 = nn.Linear(8 * 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 * 5 * 5)
+        x = x.view(-1, 8 * 4 * 4)
         # x = F.relu(self.fc1(x))
         # x = F.relu(self.fc2(x))
         # x = self.fc3(x)
@@ -51,12 +51,16 @@ class ConvNet(nn.Module):
         return x
 
     def printFector(self, x, label):
-        show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]).cpu(), "input_image.png")
+        show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]), "input_image.png", Contrast=[0, 1.0])
+        # show.DumpTensorToLog(x, "input_image.log")
         x = self.conv1(x)
         w = self.conv1.weight
-        show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv1_weight.png")
+        show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), "conv1_weight.png", Contrast=[-1.0, 1.0])
+        # show.DumpTensorToLog(w, "conv1_weight.log")
+
+        show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]), "conv1_output.png", Contrast=[-1.0, 1.0])
+        # show.DumpTensorToLog(x, "conv1_output.png")
 
-        show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]).cpu(), "conv1_output.png")
         x = self.pool(F.relu(x))
         x = self.conv2(x)
         w = self.conv2.weight
@@ -64,10 +68,10 @@ class ConvNet(nn.Module):
 
         show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]).cpu(), "conv2_output.png")
         x = self.pool(F.relu(x))
-        x = x.view(-1, 8 * 5 * 5)
+        x = x.view(-1, 8 * 4 * 4)
         x = self.fc1(x)
 
-        show.DumpTensorToImage(self.fc1.weight.view(-1, 10, 10).permute(2, 0, 1).cpu(), "fc_weight.png")
+        show.DumpTensorToImage(self.fc1.weight.view(-1, 16, 8).permute(2, 0, 1), "fc_weight.png")
 
         show.DumpTensorToImage(x.view(-1).cpu(), "fc_output.png")
 
@@ -79,8 +83,8 @@ class ConvNet(nn.Module):
         w = self.conv1.weight.grad
         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]).cpu(), "conv2_weight_grad.png")
+        show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), "conv2_weight_grad.png")
-        show.DumpTensorToImage(self.fc1.weight.grad.view(-1, 10, 10).permute(2, 0, 1).cpu(), "fc_weight_grad.png")
+        show.DumpTensorToImage(self.fc1.weight.grad.view(-1, 16, 8).permute(2, 0, 1), "fc_weight_grad.png")
 
 
 model = ConvNet().to(device)
@@ -105,9 +109,10 @@ for epoch in range(num_epochs):
         loss.backward()
         optimizer.step()
 
-        if (i + 1) % 2000 == 0:
+        if (i + 1) % 100 == 0:
             print(f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.4f}")
 
+test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False)
 for images, labels in test_loader:
     images = images.to(device)
     labels = labels.to(device)