import os import sys import torch import torch.nn as nn import torch.nn.functional as F # Add this line import torchvision import torchvision.transforms as transforms sys.path.append("..") from tools import show seed = 4321 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") num_epochs = 1 batch_size = 64 transform = transforms.Compose([transforms.ToTensor()]) train_dataset = torchvision.datasets.MNIST(root="./data", train=True, download=True, transform=transform) test_dataset = torchvision.datasets.MNIST(root="./data", train=False, download=True, transform=transform) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True) test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False) class ConvNet(nn.Module): def __init__(self): super(ConvNet, self).__init__() self.conv1 = nn.Conv2d(1, 8, 5, 1, 0) self.pool = nn.MaxPool2d(2, 2) self.conv2 = nn.Conv2d(8, 1, 5, 1, 0) self.fc1 = nn.Linear(1 * 4 * 4, 10) def forward(self, x): x = self.pool(self.conv1(x)) x = self.pool(self.conv2(x)) x = x.view(x.shape[0], -1) x = self.fc1(x) return x def forward_unsuper(self, x): x = self.conv1(x) # x = self.pool(self.conv1(x)) return x def forward_finetune(self, x): x = self.pool(self.conv1(x)) x = self.pool(self.conv2(x)) x = x.view(x.shape[0], -1) x = self.fc1(x) return x def printFector(self, x, label, dir=""): show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]), dir + "/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]), dir + "/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]), dir + "/conv1_output.png", Contrast=[-1.0, 1.0]) # show.DumpTensorToLog(x, "conv1_output.png") 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(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) x = self.fc1(x) show.DumpTensorToImage( self.fc1.weight.view(-1, pool_shape[2], pool_shape[3]), dir + "/fc_weight.png", Contrast=[-1.0, 1.0] ) show.DumpTensorToImage(x.view(-1).cpu(), dir + "/fc_output.png") criterion = nn.CrossEntropyLoss() loss = criterion(x, label) loss.backward() if self.conv1.weight.requires_grad: w = self.conv1.weight.grad show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), dir + "/conv1_weight_grad.png") if self.conv2.weight.requires_grad: w = self.conv2.weight.grad show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), dir + "/conv2_weight_grad.png") if self.fc1.weight.requires_grad: show.DumpTensorToImage( self.fc1.weight.grad.view(-1, pool_shape[2], pool_shape[3]), dir + "/fc_weight_grad.png" ) model = ConvNet().to(device) model.train() # Train the model unsuper epochs = 2 model.conv1.weight.requires_grad = True model.conv2.weight.requires_grad = False model.fc1.weight.requires_grad = False n_total_steps = len(train_loader) for epoch in range(epochs): for i, (images, labels) in enumerate(train_loader): images = images.to(device) outputs = model.forward_unsuper(images) 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 diff_ratio_mean = torch.mean(diff_ratio * diff_ratio, dim=1) label = diff_ratio_mean * 0.5 loss = F.l1_loss(diff_ratio_mean, label) model.conv1.weight.grad = None loss.backward() 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}") # Train the model model.conv1.weight.requires_grad = False model.conv2.weight.requires_grad = True model.fc1.weight.requires_grad = True criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.2) n_total_steps = len(train_loader) for epoch in range(num_epochs): for i, (images, labels) in enumerate(train_loader): images = images.to(device) labels = labels.to(device) outputs = model.forward_finetune(images) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() 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") test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False) test_loader = iter(test_loader) images, labels = next(test_loader) images = images.to(device) labels = labels.to(device) model.printFector(images, labels, "dump1") images, labels = next(test_loader) images = images.to(device) labels = labels.to(device) model.printFector(images, labels, "dump2") # Test the model with torch.no_grad(): n_correct = 0 n_samples = 0 for images, labels in test_loader: images = images.to(device) labels = labels.to(device) outputs = model(images) # max returns (value ,index) _, predicted = torch.max(outputs.data, 1) n_samples += labels.size(0) n_correct += (predicted == labels).sum().item() acc = 100.0 * n_correct / n_samples print(f"Accuracy of the network on the 10000 test images: {acc} %")