From 81f203ce5984821987c79a2973c88605e4919d1d Mon Sep 17 00:00:00 2001
From: Colin <colin>
Date: Sun, 22 Sep 2024 15:18:08 +0800
Subject: [PATCH] update the max grad's weight.

---
 unsuper/minist.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/unsuper/minist.py b/unsuper/minist.py
index ec02507..52bb775 100644
--- a/unsuper/minist.py
+++ b/unsuper/minist.py
@@ -65,17 +65,16 @@ class ConvNet(nn.Module):
         show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]), dir + "/conv1_output.png", Contrast=[-1.0, 1.0])
         # show.DumpTensorToLog(x, "conv1_output.png")
 
-        x = self.pool(F.relu(x))
+        x = self.pool(x)
         x = self.conv2(x)
         w = self.conv2.weight
         show.DumpTensorToImage(
             w.view(-1, w.shape[2], w.shape[3]).cpu(), dir + "/conv2_weight.png", Contrast=[-1.0, 1.0]
         )
-
         show.DumpTensorToImage(
             x.view(-1, x.shape[2], x.shape[3]).cpu(), dir + "/conv2_output.png", Contrast=[-1.0, 1.0]
         )
-        x = self.pool(F.relu(x))
+        x = self.pool(x)
         show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]).cpu(), dir + "/pool_output.png", Contrast=[-1.0, 1.0])
         pool_shape = x.shape
         x = x.view(x.shape[0], -1)
@@ -105,7 +104,7 @@ model = ConvNet().to(device)
 model.train()
 
 # Train the model unsuper
-epochs = 10
+epochs = 2
 model.conv1.weight.requires_grad = True
 model.conv2.weight.requires_grad = False
 model.fc1.weight.requires_grad = False
@@ -114,7 +113,10 @@ for epoch in range(epochs):
     for i, (images, labels) in enumerate(train_loader):
         images = images.to(device)
         outputs = model.forward_unsuper(images)
-        sample = outputs.view(outputs.shape[0], -1)
+
+        outputs = outputs.permute(1, 0, 2, 3)  # 64 8 24 24 -> 8 64 24 24
+        sample = outputs.reshape(outputs.shape[0], -1)  # -> 8 36864
+
         sample_mean = torch.mean(sample, dim=1, keepdim=True)
         diff_mean = torch.mean(torch.abs(sample - sample_mean), dim=1, keepdim=True)
         diff_ratio = (sample - sample_mean) / diff_mean
@@ -123,7 +125,13 @@ for epoch in range(epochs):
         loss = F.l1_loss(diff_ratio_mean, label)
         model.conv1.weight.grad = None
         loss.backward()
-        model.conv1.weight.data = model.conv1.weight.data - model.conv1.weight.grad * 0.2
+
+        grad = model.conv1.weight.data
+        grad = grad.view(8, -1)
+        grad_mean = torch.mean(grad, dim=1)
+        max, index = torch.max(grad_mean, dim=0)
+
+        model.conv1.weight.data[index] = model.conv1.weight.data[index] - model.conv1.weight.grad[index] * 0.2
         if (i + 1) % 100 == 0:
             print(f"Epoch [{epoch+1}/{epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.8f}")