From f3690fd47fe496706fd12d6c266a2b49d2d88358 Mon Sep 17 00:00:00 2001
From: Colin <colin>
Date: Mon, 7 Oct 2024 16:39:29 +0800
Subject: [PATCH] Update unsuper learning.

---
 unsuper/minist.py | 26 +++++++++++++++++++-------
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/unsuper/minist.py b/unsuper/minist.py
index c409962..73818a3 100644
--- a/unsuper/minist.py
+++ b/unsuper/minist.py
@@ -105,7 +105,7 @@ model = ConvNet().to(device)
 model.train()
 
 # Train the model unsuper
-epochs = 10
+epochs = 20
 model.conv1.weight.requires_grad = True
 model.conv2.weight.requires_grad = False
 model.fc1.weight.requires_grad = False
@@ -117,16 +117,24 @@ for epoch in range(epochs):
 
         outputs = outputs.permute(0, 2, 3, 1)  # 64 8 24 24 -> 64 24 24 8
         sample = outputs.reshape(-1, outputs.shape[3])  # -> 36864 8
-        abs = torch.abs(sample)
+        abs = torch.abs(sample).detach()
         max, max_index = torch.max(abs, dim=1)
-        label = sample * 0.9
-        all = range(0, label.shape[0])
-        label[all, max_index] = label[all, max_index] * 1.1
+        mean = torch.mean(abs, dim=1)
+        mean = torch.expand_copy(mean.reshape(-1, 1), sample.shape)
+        max = torch.expand_copy(max.reshape(-1, 1), sample.shape)
+
+        all = range(0, sample.shape[0])
+        ratio_max = abs / mean
+        ratio_nor = (max - abs) / max
+        ratio_nor[all, max_index] = ratio_max[all, max_index].clone()
+        ratio_nor = torch.where(torch.isnan(ratio_nor), 1.0, ratio_nor)
+        label = sample * ratio_nor
+
         loss = F.l1_loss(sample, label)
         model.conv1.weight.grad = None
         loss.backward()
 
-        model.conv1.weight.data = model.conv1.weight.data - model.conv1.weight.grad * 100
+        model.conv1.weight.data = model.conv1.weight.data - model.conv1.weight.grad * 10
 
         if (i + 1) % 100 == 0:
             print(f"Epoch [{epoch+1}/{epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.8f}")
@@ -136,6 +144,10 @@ show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv1_weight_g
 w = model.conv1.weight.data
 show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), "conv1_weight_update.png", Contrast=[-1.0, 1.0])
 
+# loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False)
+# images, labels = next(iter(loader))
+# images = images.to(device)
+
 # Train the model
 model.conv1.weight.requires_grad = False
 model.conv2.weight.requires_grad = True
@@ -155,7 +167,7 @@ for epoch in range(num_epochs):
         if (i + 1) % 100 == 0:
             print(f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.4f}")
 
-# print("Finished Training")
+print("Finished Training")
 
 test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False)
 test_loader = iter(test_loader)