colin
/
Witllm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update the max grad's weight.

This commit is contained in:
Colin 2024-09-22 15:18:08 +08:00
parent 03516f6302
commit 81f203ce59
1 changed files with 14 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
unsuper/minist.py
View File

@ -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}")