MainCode.py

# -*- coding: utf-8 -*-
"""MJAhmadi_NNDL_HW2_Q1.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1YPqjfZnfNNE9yv4msI5fQ1bvv4TlpKeP

# 1.1
"""

import torch
import torchvision
import torchvision.transforms as transforms
import random
import os

# set random seed
random_seed = 42
torch.manual_seed(random_seed)
random.seed(random_seed)

# download datasets
main_trainset_cifar10 = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=None)
train_mean_cifar10 = main_trainset_cifar10.data.mean(axis=(0, 1, 2))/255
train_std_cifar10 = main_trainset_cifar10.data.std(axis=(0, 1, 2))/255

main_testset_cifar10 = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=None)
test_mean_cifar10 = main_testset_cifar10.data.mean(axis=(0, 1, 2))/255
test_std_cifar10 = main_testset_cifar10.data.std(axis=(0, 1, 2))/255

main_trainset_mnist = torchvision.datasets.MNIST(root='./data', train=True,
                                        download=True, transform=None)
train_mean_mnist = main_trainset_mnist.data.float().mean()/255
train_std_mnist = main_trainset_mnist.data.float().std()/255

main_testset_mnist = torchvision.datasets.MNIST(root='./data', train=False,
                                       download=True, transform=None)
test_mean_mnist = main_testset_mnist.data.float().mean()/255
test_std_mnist = main_testset_mnist.data.float().std()/255

main_trainset_fashion_mnist = torchvision.datasets.FashionMNIST(root='./data', train=True,
                                        download=True, transform=None)
train_mean_fashion_mnist = main_trainset_fashion_mnist.data.float().mean()/255
train_std_fashion_mnist = main_trainset_fashion_mnist.data.float().std()/255

main_testset_fashion_mnist = torchvision.datasets.FashionMNIST(root='./data', train=False,
                                       download=True, transform=None)
test_mean_fashion_mnist = main_testset_fashion_mnist.data.float().mean()/255
test_std_fashion_mnist = main_testset_fashion_mnist.data.float().std()/255

# define data transformations
transform_cifar10 = transforms.Compose(
    [transforms.Grayscale(num_output_channels=1),
     transforms.ToTensor(),
     transforms.Normalize(mean=train_mean_cifar10, std=train_std_cifar10),
     transforms.Resize((28, 28))]
)

# define data transformations
transform_mnist = transforms.Compose(
    [transforms.Grayscale(num_output_channels=1),
     transforms.ToTensor(),
     transforms.Normalize(mean=train_mean_mnist, std=train_std_mnist),
     transforms.Resize((28, 28))]
)

# define flip transform for Fashion-MNIST dataset
transform_fashion_mnist = transforms.Compose(
    [transforms.RandomHorizontalFlip(p=0.5),
     transforms.Grayscale(num_output_channels=1),
     transforms.ToTensor(),
     transforms.Normalize(mean=train_mean_fashion_mnist, std=train_std_fashion_mnist),
     transforms.Resize((28, 28))]
)

# download datasets
trainset_cifar10 = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform_cifar10)
testset_cifar10 = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform_cifar10)

trainset_mnist = torchvision.datasets.MNIST(root='./data', train=True,
                                        download=True, transform=transform_mnist)
testset_mnist = torchvision.datasets.MNIST(root='./data', train=False,
                                       download=True, transform=transform_mnist)

trainset_fashion_mnist = torchvision.datasets.FashionMNIST(root='./data', train=True,
                                        download=True, transform=transform_fashion_mnist)
testset_fashion_mnist = torchvision.datasets.FashionMNIST(root='./data', train=False,
                                       download=True, transform=transform_fashion_mnist)

# split trainset into train and validation sets
trainset_cifar10, valset_cifar10 = torch.utils.data.random_split(trainset_cifar10, [40000, 10000])
trainset_mnist, valset_mnist = torch.utils.data.random_split(trainset_mnist, [50000, 10000])
trainset_fashion_mnist, valset_fashion_mnist = torch.utils.data.random_split(trainset_fashion_mnist, [48000, 12000])

# save datasets
os.makedirs('./data/preprocessed/', exist_ok=True)
torch.save(trainset_cifar10, './data/preprocessed/trainset_cifar10.pt')
torch.save(valset_cifar10, './data/preprocessed/valset_cifar10.pt')
torch.save(testset_cifar10, './data/preprocessed/testset_cifar10.pt')

torch.save(trainset_mnist, './data/preprocessed/trainset_mnist.pt')
torch.save(valset_mnist, './data/preprocessed/valset_mnist.pt')
torch.save(testset_mnist, './data/preprocessed/testset_mnist.pt')

torch.save(trainset_fashion_mnist, './data/preprocessed/trainset_fashion_mnist.pt')
torch.save(valset_fashion_mnist, './data/preprocessed/valset_fashion_mnist.pt')
torch.save(testset_fashion_mnist, './data/preprocessed/testset_fashion_mnist.pt')

print("CIFAR10 mean (Train): ", train_mean_cifar10)
print("CIFAR10 std (Train): ", train_std_cifar10)
print("CIFAR10 mean (Test): ", test_mean_cifar10)
print("CIFAR10 std (Test): ", test_std_cifar10)

print("MNIST mean (Train): ", train_mean_mnist)
print("MNIST std (Train): ", train_std_mnist)
print("MNIST mean (Test): ", test_mean_mnist)
print("MNIST std (Test): ", test_std_mnist)

print("FASHION-MNIST mean (Train): ", train_mean_fashion_mnist)
print("FASHION-MNIST std (Train): ", train_std_fashion_mnist)
print("FASHION-MNIST mean (Test): ", test_mean_fashion_mnist)
print("FASHION-MNIST std (Test): ", test_std_fashion_mnist)

import torch
import torchvision
import torchvision.transforms as transforms
import random
import os

# set random seed
random_seed = 42
torch.manual_seed(random_seed)
random.seed(random_seed)

# define data transformations
transform = transforms.Compose(
    [transforms.Grayscale(num_output_channels=1),
     transforms.ToTensor(),
     transforms.Normalize((0.5,), (0.5,)),
     transforms.Resize((28, 28))]
)

# define flip transform for Fashion-MNIST dataset
flip_transform = transforms.Compose(
    [transforms.RandomHorizontalFlip(p=0.5),
     transforms.Grayscale(num_output_channels=1),
     transforms.ToTensor(),
     transforms.Normalize((0.5,), (0.5,)),
     transforms.Resize((28, 28))]
)

# download datasets
trainset_cifar10 = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
testset_cifar10 = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)

trainset_mnist = torchvision.datasets.MNIST(root='./data', train=True,
                                        download=True, transform=transform)
testset_mnist = torchvision.datasets.MNIST(root='./data', train=False,
                                       download=True, transform=transform)

trainset_fashion_mnist = torchvision.datasets.FashionMNIST(root='./data', train=True,
                                        download=True, transform=flip_transform)
testset_fashion_mnist = torchvision.datasets.FashionMNIST(root='./data', train=False,
                                       download=True, transform=flip_transform)

# split trainset into train and validation sets
trainset_cifar10, valset_cifar10 = torch.utils.data.random_split(trainset_cifar10, [40000, 10000])
trainset_mnist, valset_mnist = torch.utils.data.random_split(trainset_mnist, [50000, 10000])
trainset_fashion_mnist, valset_fashion_mnist = torch.utils.data.random_split(trainset_fashion_mnist, [48000, 12000])

# save datasets
os.makedirs('./data/preprocessed/', exist_ok=True)
torch.save(trainset_cifar10, './data/preprocessed/trainset_cifar10.pt')
torch.save(valset_cifar10, './data/preprocessed/valset_cifar10.pt')
torch.save(testset_cifar10, './data/preprocessed/testset_cifar10.pt')

torch.save(trainset_mnist, './data/preprocessed/trainset_mnist.pt')
torch.save(valset_mnist, './data/preprocessed/valset_mnist.pt')
torch.save(testset_mnist, './data/preprocessed/testset_mnist.pt')

torch.save(trainset_fashion_mnist, './data/preprocessed/trainset_fashion_mnist.pt')
torch.save(valset_fashion_mnist, './data/preprocessed/valset_fashion_mnist.pt')
torch.save(testset_fashion_mnist, './data/preprocessed/testset_fashion_mnist.pt')

"""# 1.2

## MNIST
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# Load the preprocessed datasets
trainset_mnist = torch.load('./data/preprocessed/trainset_mnist.pt')
valset_mnist = torch.load('./data/preprocessed/valset_mnist.pt')
testset_mnist = torch.load('./data/preprocessed/testset_mnist.pt')

# Define hyperparameters
batch_size = 128
learning_rate = 0.02
momentum = 0.9
weight_decay = 0.000005
num_epochs = 150

# Create data loaders
trainloader_mnist = DataLoader(trainset_mnist, batch_size=batch_size, shuffle=True)
valloader_mnist = DataLoader(valset_mnist, batch_size=batch_size, shuffle=False)
testloader_mnist = DataLoader(testset_mnist, batch_size=batch_size, shuffle=False)

# Define the model
class SCNNB(nn.Module):
    def __init__(self):
        super(SCNNB, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-a
class SCNNB_a(nn.Module):
    def __init__(self):
        super(SCNNB_a, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-b
class SCNNB_b(nn.Module):
    def __init__(self):
        super(SCNNB_b, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# define the models
models = [SCNNB(), SCNNB_a(), SCNNB_b()]
colors = ['blue', 'green', 'red']

# define the training loop
train_losses = [[] for _ in range(len(models))]
train_accs = [[] for _ in range(len(models))]
val_losses = [[] for _ in range(len(models))]
val_accs = [[] for _ in range(len(models))]
test_accs = [[] for _ in range(len(models))]

for i, model in enumerate(models):
    print(f"Training model {i+1}")
    for epoch in range(num_epochs):

        criterion = nn.CrossEntropyLoss()
        optimizer = optim.SGD(model.parameters(), lr=0.02, momentum=0.9, weight_decay=0.000005)
        
        train_loss = 0.0
        train_total = 0
        train_correct = 0
        val_loss = 0.0
        val_total = 0
        val_correct = 0
        train_acc = 0.0
        val_acc = 0.0

        # train the model on training set
        model.train()
        for j, data in enumerate(trainloader_mnist, 0):
            inputs, labels = data
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            train_acc += (predicted == labels).sum().item()

        # evaluate the model on validation set
        model.eval()
        with torch.no_grad():
            for j, data in enumerate(valloader_mnist, 0):
                inputs, labels = data
                outputs = model(inputs)
                loss = criterion(outputs, labels)

                val_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                val_acc += (predicted == labels).sum().item()

        # calculate average loss and accuracy for training and validation sets
        train_loss /= len(trainloader_mnist)
        train_acc /= len(trainset_mnist)
        val_loss /= len(valloader_mnist)
        val_acc /= len(valset_mnist)
        train_losses[i].append(train_loss)
        train_accs[i].append(train_acc)
        val_losses[i].append(val_loss)
        val_accs[i].append(val_acc)

        # evaluate the model on test set
        model.eval()
        with torch.no_grad():
            test_acc = 0.0
            for j, data in enumerate(testloader_mnist, 0):
                inputs, labels = data
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                test_acc += (predicted == labels).sum().item()

            test_acc /= len(testset_mnist)
            test_accs[i].append(test_acc)

        # print the loss and accuracy for each epoch
        print(f'Model {i+1} | Epoch {epoch+1:3d} | Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f} | Test Acc: {test_acc:.4f}')


# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(train_losses[i], color=colors[i], label=f'Train Loss ({models[i].__class__.__name__})')
    plt.plot(val_losses[i], '--', color=colors[i], label=f'Val Loss ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('Training and Validation Loss over Epochs')
plt.savefig("Q21figure1.pdf")
plt.show()


# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(train_accs[i], color=colors[i], label=f'Train Acc ({models[i].__class__.__name__})')
    plt.plot(val_accs[i], '--', color=colors[i], label=f'Val Acc ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.title('Training and Validation Loss over Epochs')
plt.savefig("Q21figure2.pdf")
plt.show()

# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(test_accs[i], color=colors[i], label=f'Test Acc ({models[i].__class__.__name__})')
    # plt.plot(val_accs[i], '--', color=colors[i], label=f'Val Loss ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('Test Accuracy over Epochs')
plt.savefig("Q21figure3.pdf")
plt.show()

# evaluate the model on test set
for i, model in enumerate(models):
    print(f"Test model {i+1}")
    model.eval()
    with torch.no_grad():
      test_acc = 0.0
      all_targets = []
      all_predictions = []
      for i, data in enumerate(testloader_mnist, 0):
          inputs, labels = data
          outputs = model(inputs)
          _, predicted = torch.max(outputs.data, 1)
          test_acc += (predicted == labels).sum().item()
          all_targets.extend(labels.cpu().numpy())
          all_predictions.extend(predicted.cpu().numpy())

      # calculate accuracy on test set
      test_acc /= len(testset_mnist)
      test_accs.append(test_acc)
      
      # print classification report
      print('Classification report:')
      print(classification_report(all_targets, all_predictions))
      
      # plot confusion matrix
      for i in range(len(models)):
        cm = confusion_matrix(all_targets, all_predictions)
        ax = sns.heatmap(cm, annot=True, fmt='g', cmap='Blues', cbar=False)
        ax.set_xlabel('Predicted labels')
        ax.set_ylabel('True labels')
        ax.xaxis.set_ticklabels(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])
        ax.yaxis.set_ticklabels(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])
        plt.title(f'Confusion Matrix ({models[i].__class__.__name__})')
        plt.savefig(f"Q21figure4_{models[i].__class__.__name__}.pdf")
        plt.show()

# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(train_accs[i], color=colors[i], label=f'Train Acc ({models[i].__class__.__name__})')
    plt.plot(val_accs[i], '--', color=colors[i], label=f'Val Acc ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.title('Training and Validation Accuracy over Epochs')
plt.savefig("Q21figure2.pdf")
plt.show()

# evaluate the model on test set
for i, model in enumerate(models):
    print(f"Test model {i+1}")
    model.eval()
    with torch.no_grad():
      test_acc = 0.0
      all_targets = []
      all_predictions = []
      for i, data in enumerate(testloader_mnist, 0):
          inputs, labels = data
          outputs = model(inputs)
          _, predicted = torch.max(outputs.data, 1)
          test_acc += (predicted == labels).sum().item()
          all_targets.extend(labels.cpu().numpy())
          all_predictions.extend(predicted.cpu().numpy())

      # calculate accuracy on test set
      test_acc /= len(testset_mnist)
      test_accs.append(test_acc)
      
      # print classification report
      print('Classification report:')
      print(classification_report(all_targets, all_predictions))
      
      # plot confusion matrix
      cm = confusion_matrix(all_targets, all_predictions)
      ax = sns.heatmap(cm, annot=True, fmt='g', cmap='Blues', cbar=False)
      ax.set_xlabel('Predicted labels')
      ax.set_ylabel('True labels')
      ax.xaxis.set_ticklabels(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])
      ax.yaxis.set_ticklabels(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])
      plt.title(f'Confusion Matrix ({model.__class__.__name__})')
      plt.savefig(f"confusion_matrix_{model.__class__.__name__}.pdf")
      plt.show()

"""## Fashion-MNIST"""



import torch
import torchvision
import torchvision.transforms as transforms
import random
import os

# set random seed
random_seed = 42
torch.manual_seed(random_seed)
random.seed(random_seed)

# download datasets
main_trainset_cifar10 = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=None)
train_mean_cifar10 = main_trainset_cifar10.data.mean(axis=(0, 1, 2))/255
train_std_cifar10 = main_trainset_cifar10.data.std(axis=(0, 1, 2))/255

main_testset_cifar10 = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=None)
test_mean_cifar10 = main_testset_cifar10.data.mean(axis=(0, 1, 2))/255
test_std_cifar10 = main_testset_cifar10.data.std(axis=(0, 1, 2))/255

main_trainset_mnist = torchvision.datasets.MNIST(root='./data', train=True,
                                        download=True, transform=None)
train_mean_mnist = main_trainset_mnist.data.float().mean()/255
train_std_mnist = main_trainset_mnist.data.float().std()/255

main_testset_mnist = torchvision.datasets.MNIST(root='./data', train=False,
                                       download=True, transform=None)
test_mean_mnist = main_testset_mnist.data.float().mean()/255
test_std_mnist = main_testset_mnist.data.float().std()/255

main_trainset_fashion_mnist = torchvision.datasets.FashionMNIST(root='./data', train=True,
                                        download=True, transform=None)
train_mean_fashion_mnist = main_trainset_fashion_mnist.data.float().mean()/255
train_std_fashion_mnist = main_trainset_fashion_mnist.data.float().std()/255

main_testset_fashion_mnist = torchvision.datasets.FashionMNIST(root='./data', train=False,
                                       download=True, transform=None)
test_mean_fashion_mnist = main_testset_fashion_mnist.data.float().mean()/255
test_std_fashion_mnist = main_testset_fashion_mnist.data.float().std()/255

# define data transformations
transform_cifar10 = transforms.Compose(
    [transforms.Grayscale(num_output_channels=1),
     transforms.ToTensor(),
     transforms.Normalize(mean=train_mean_cifar10, std=train_std_cifar10),
     transforms.Resize((28, 28))]
)

# define data transformations
transform_mnist = transforms.Compose(
    [transforms.Grayscale(num_output_channels=1),
     transforms.ToTensor(),
     transforms.Normalize(mean=train_mean_mnist, std=train_std_mnist),
     transforms.Resize((28, 28))]
)

# define flip transform for Fashion-MNIST dataset
transform_fashion_mnist = transforms.Compose(
    [transforms.RandomHorizontalFlip(p=0.5),
     transforms.Grayscale(num_output_channels=1),
     transforms.ToTensor(),
     transforms.Normalize(mean=train_mean_fashion_mnist, std=train_std_fashion_mnist),
     transforms.Resize((28, 28))]
)

# download datasets
trainset_cifar10 = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform_cifar10)
testset_cifar10 = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform_cifar10)

trainset_mnist = torchvision.datasets.MNIST(root='./data', train=True,
                                        download=True, transform=transform_mnist)
testset_mnist = torchvision.datasets.MNIST(root='./data', train=False,
                                       download=True, transform=transform_mnist)

trainset_fashion_mnist = torchvision.datasets.FashionMNIST(root='./data', train=True,
                                        download=True, transform=transform_fashion_mnist)
testset_fashion_mnist = torchvision.datasets.FashionMNIST(root='./data', train=False,
                                       download=True, transform=transform_fashion_mnist)

# split trainset into train and validation sets
trainset_cifar10, valset_cifar10 = torch.utils.data.random_split(trainset_cifar10, [40000, 10000])
trainset_mnist, valset_mnist = torch.utils.data.random_split(trainset_mnist, [50000, 10000])
trainset_fashion_mnist, valset_fashion_mnist = torch.utils.data.random_split(trainset_fashion_mnist, [48000, 12000])

# save datasets
os.makedirs('./data/preprocessed/', exist_ok=True)
torch.save(trainset_cifar10, './data/preprocessed/trainset_cifar10.pt')
torch.save(valset_cifar10, './data/preprocessed/valset_cifar10.pt')
torch.save(testset_cifar10, './data/preprocessed/testset_cifar10.pt')

torch.save(trainset_mnist, './data/preprocessed/trainset_mnist.pt')
torch.save(valset_mnist, './data/preprocessed/valset_mnist.pt')
torch.save(testset_mnist, './data/preprocessed/testset_mnist.pt')

torch.save(trainset_fashion_mnist, './data/preprocessed/trainset_fashion_mnist.pt')
torch.save(valset_fashion_mnist, './data/preprocessed/valset_fashion_mnist.pt')
torch.save(testset_fashion_mnist, './data/preprocessed/testset_fashion_mnist.pt')

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# Load the preprocessed datasets
trainset_mnist = torch.load('./data/preprocessed/trainset_fashion_mnist.pt')
valset_mnist = torch.load('./data/preprocessed/valset_fashion_mnist.pt')
testset_mnist = torch.load('./data/preprocessed/testset_fashion_mnist.pt')

# Define hyperparameters
batch_size = 128
learning_rate = 0.02
momentum = 0.9
weight_decay = 0.000005
num_epochs = 150

# Create data loaders
trainloader_mnist = DataLoader(trainset_mnist, batch_size=batch_size, shuffle=True)
valloader_mnist = DataLoader(valset_mnist, batch_size=batch_size, shuffle=False)
testloader_mnist = DataLoader(testset_mnist, batch_size=batch_size, shuffle=False)

# Define the model
class SCNNB(nn.Module):
    def __init__(self):
        super(SCNNB, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-a
class SCNNB_a(nn.Module):
    def __init__(self):
        super(SCNNB_a, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-b
class SCNNB_b(nn.Module):
    def __init__(self):
        super(SCNNB_b, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# define the models
models = [SCNNB(), SCNNB_a(), SCNNB_b()]
colors = ['blue', 'green', 'red']

# define the training loop
train_losses = [[] for _ in range(len(models))]
train_accs = [[] for _ in range(len(models))]
val_losses = [[] for _ in range(len(models))]
val_accs = [[] for _ in range(len(models))]
test_accs = [[] for _ in range(len(models))]

for i, model in enumerate(models):
    print(f"Training model {i+1}")
    for epoch in range(num_epochs):

        criterion = nn.CrossEntropyLoss()
        optimizer = optim.SGD(model.parameters(), lr=0.02, momentum=0.9, weight_decay=0.000005)
        
        train_loss = 0.0
        train_total = 0
        train_correct = 0
        val_loss = 0.0
        val_total = 0
        val_correct = 0
        train_acc = 0.0
        val_acc = 0.0

        # train the model on training set
        model.train()
        for j, data in enumerate(trainloader_mnist, 0):
            inputs, labels = data
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            train_acc += (predicted == labels).sum().item()

        # evaluate the model on validation set
        model.eval()
        with torch.no_grad():
            for j, data in enumerate(valloader_mnist, 0):
                inputs, labels = data
                outputs = model(inputs)
                loss = criterion(outputs, labels)

                val_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                val_acc += (predicted == labels).sum().item()

        # calculate average loss and accuracy for training and validation sets
        train_loss /= len(trainloader_mnist)
        train_acc /= len(trainset_mnist)
        val_loss /= len(valloader_mnist)
        val_acc /= len(valset_mnist)
        train_losses[i].append(train_loss)
        train_accs[i].append(train_acc)
        val_losses[i].append(val_loss)
        val_accs[i].append(val_acc)

        # evaluate the model on test set
        model.eval()
        with torch.no_grad():
            test_acc = 0.0
            for j, data in enumerate(testloader_mnist, 0):
                inputs, labels = data
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                test_acc += (predicted == labels).sum().item()

            test_acc /= len(testset_mnist)
            test_accs[i].append(test_acc)

        # print the loss and accuracy for each epoch
        print(f'Model {i+1} | Epoch {epoch+1:3d} | Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f} | Test Acc: {test_acc:.4f}')


# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(train_losses[i], color=colors[i], label=f'Train Loss ({models[i].__class__.__name__})')
    plt.plot(val_losses[i], '--', color=colors[i], label=f'Val Loss ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('Training and Validation Loss over Epochs')
plt.savefig("Q21figure1.pdf")
plt.show()


# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(train_accs[i], color=colors[i], label=f'Train Acc ({models[i].__class__.__name__})')
    plt.plot(val_accs[i], '--', color=colors[i], label=f'Val Acc ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.title('Training and Validation Accuracy over Epochs')
plt.savefig("Q21figure2.pdf")
plt.show()

# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(test_accs[i], color=colors[i], label=f'Test Acc ({models[i].__class__.__name__})')
    # plt.plot(val_accs[i], '--', color=colors[i], label=f'Val Loss ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('Test Accuracy over Epochs')
plt.savefig("Q21figure3.pdf")
plt.show()

# evaluate the model on test set
for i, model in enumerate(models):
    print(f"Test model {i+1}")
    model.eval()
    with torch.no_grad():
      test_acc = 0.0
      all_targets = []
      all_predictions = []
      for i, data in enumerate(testloader_mnist, 0):
          inputs, labels = data
          outputs = model(inputs)
          _, predicted = torch.max(outputs.data, 1)
          test_acc += (predicted == labels).sum().item()
          all_targets.extend(labels.cpu().numpy())
          all_predictions.extend(predicted.cpu().numpy())

      # calculate accuracy on test set
      test_acc /= len(testset_mnist)
      test_accs.append(test_acc)
      
      # print classification report
      print('Classification report:')
      print(classification_report(all_targets, all_predictions))
      
      # plot confusion matrix
      cm = confusion_matrix(all_targets, all_predictions)
      ax = sns.heatmap(cm, annot=True, fmt='g', cmap='Blues', cbar=False)
      ax.set_xlabel('Predicted labels')
      ax.set_ylabel('True labels')
      ax.xaxis.set_ticklabels(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])
      ax.yaxis.set_ticklabels(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])
      plt.title(f'Confusion Matrix ({model.__class__.__name__})')
      plt.savefig(f"confusion_matrix_{model.__class__.__name__}.pdf")
      plt.show()

"""## CIFAR10"""

import torch
import torchvision
import torchvision.transforms as transforms
import random
import os

# set random seed
random_seed = 42
torch.manual_seed(random_seed)
random.seed(random_seed)

# download datasets
main_trainset_cifar10 = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=None)
train_mean_cifar10 = main_trainset_cifar10.data.mean(axis=(0, 1, 2))/255
train_std_cifar10 = main_trainset_cifar10.data.std(axis=(0, 1, 2))/255

main_testset_cifar10 = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=None)
test_mean_cifar10 = main_testset_cifar10.data.mean(axis=(0, 1, 2))/255
test_std_cifar10 = main_testset_cifar10.data.std(axis=(0, 1, 2))/255

main_trainset_mnist = torchvision.datasets.MNIST(root='./data', train=True,
                                        download=True, transform=None)
train_mean_mnist = main_trainset_mnist.data.float().mean()/255
train_std_mnist = main_trainset_mnist.data.float().std()/255

main_testset_mnist = torchvision.datasets.MNIST(root='./data', train=False,
                                       download=True, transform=None)
test_mean_mnist = main_testset_mnist.data.float().mean()/255
test_std_mnist = main_testset_mnist.data.float().std()/255

main_trainset_fashion_mnist = torchvision.datasets.FashionMNIST(root='./data', train=True,
                                        download=True, transform=None)
train_mean_fashion_mnist = main_trainset_fashion_mnist.data.float().mean()/255
train_std_fashion_mnist = main_trainset_fashion_mnist.data.float().std()/255

main_testset_fashion_mnist = torchvision.datasets.FashionMNIST(root='./data', train=False,
                                       download=True, transform=None)
test_mean_fashion_mnist = main_testset_fashion_mnist.data.float().mean()/255
test_std_fashion_mnist = main_testset_fashion_mnist.data.float().std()/255

# define data transformations
transform_cifar10 = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize(mean=train_mean_cifar10, std=train_std_cifar10),
     transforms.Grayscale(num_output_channels=1),
     transforms.Resize((28, 28))]
)

# define data transformations
transform_mnist = transforms.Compose(
    [transforms.Grayscale(num_output_channels=1),
     transforms.ToTensor(),
     transforms.Normalize(mean=train_mean_mnist, std=train_std_mnist),
     transforms.Resize((28, 28))]
)

# define flip transform for Fashion-MNIST dataset
transform_fashion_mnist = transforms.Compose(
    [transforms.RandomHorizontalFlip(p=0.5),
     transforms.Grayscale(num_output_channels=1),
     transforms.ToTensor(),
     transforms.Normalize(mean=train_mean_fashion_mnist, std=train_std_fashion_mnist),
     transforms.Resize((28, 28))]
)

# download datasets
trainset_cifar10 = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform_cifar10)
testset_cifar10 = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform_cifar10)

trainset_mnist = torchvision.datasets.MNIST(root='./data', train=True,
                                        download=True, transform=transform_mnist)
testset_mnist = torchvision.datasets.MNIST(root='./data', train=False,
                                       download=True, transform=transform_mnist)

trainset_fashion_mnist = torchvision.datasets.FashionMNIST(root='./data', train=True,
                                        download=True, transform=transform_fashion_mnist)
testset_fashion_mnist = torchvision.datasets.FashionMNIST(root='./data', train=False,
                                       download=True, transform=transform_fashion_mnist)

# split trainset into train and validation sets
trainset_cifar10, valset_cifar10 = torch.utils.data.random_split(trainset_cifar10, [40000, 10000])
trainset_mnist, valset_mnist = torch.utils.data.random_split(trainset_mnist, [50000, 10000])
trainset_fashion_mnist, valset_fashion_mnist = torch.utils.data.random_split(trainset_fashion_mnist, [48000, 12000])

# save datasets
os.makedirs('./data/preprocessed/', exist_ok=True)
torch.save(trainset_cifar10, './data/preprocessed/trainset_cifar10.pt')
torch.save(valset_cifar10, './data/preprocessed/valset_cifar10.pt')
torch.save(testset_cifar10, './data/preprocessed/testset_cifar10.pt')

torch.save(trainset_mnist, './data/preprocessed/trainset_mnist.pt')
torch.save(valset_mnist, './data/preprocessed/valset_mnist.pt')
torch.save(testset_mnist, './data/preprocessed/testset_mnist.pt')

torch.save(trainset_fashion_mnist, './data/preprocessed/trainset_fashion_mnist.pt')
torch.save(valset_fashion_mnist, './data/preprocessed/valset_fashion_mnist.pt')
torch.save(testset_fashion_mnist, './data/preprocessed/testset_fashion_mnist.pt')

"""### Test1"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# Load the preprocessed datasets
trainset_mnist = torch.load('./data/preprocessed/trainset_cifar10.pt')
valset_mnist = torch.load('./data/preprocessed/valset_cifar10.pt')
testset_mnist = torch.load('./data/preprocessed/testset_cifar10.pt')

# Define hyperparameters
batch_size = 128
learning_rate = 0.02
momentum = 0.9
weight_decay = 0.000005
num_epochs = 150

# Create data loaders
trainloader_mnist = DataLoader(trainset_mnist, batch_size=batch_size, shuffle=True)
valloader_mnist = DataLoader(valset_mnist, batch_size=batch_size, shuffle=False)
testloader_mnist = DataLoader(testset_mnist, batch_size=batch_size, shuffle=False)

# Define the model
class SCNNB(nn.Module):
    def __init__(self):
        super(SCNNB, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-a
class SCNNB_a(nn.Module):
    def __init__(self):
        super(SCNNB_a, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-b
class SCNNB_b(nn.Module):
    def __init__(self):
        super(SCNNB_b, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# define the models
models = [SCNNB(), SCNNB_a(), SCNNB_b()]
colors = ['blue', 'green', 'red']

# define the training loop
train_losses = [[] for _ in range(len(models))]
train_accs = [[] for _ in range(len(models))]
val_losses = [[] for _ in range(len(models))]
val_accs = [[] for _ in range(len(models))]
test_accs = [[] for _ in range(len(models))]

for i, model in enumerate(models):
    print(f"Training model {i+1}")
    for epoch in range(num_epochs):

        criterion = nn.CrossEntropyLoss()
        optimizer = optim.SGD(model.parameters(), lr=0.02, momentum=0.9, weight_decay=0.000005)
        train_loss = 0.0
        train_total = 0
        train_correct = 0
        val_loss = 0.0
        val_total = 0
        val_correct = 0
        train_acc = 0.0
        val_acc = 0.0

        # train the model on training set
        model.train()
        for j, data in enumerate(trainloader_mnist, 0):
            inputs, labels = data
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            train_acc += (predicted == labels).sum().item()

        # evaluate the model on validation set
        model.eval()
        with torch.no_grad():
            for j, data in enumerate(valloader_mnist, 0):
                inputs, labels = data
                outputs = model(inputs)
                loss = criterion(outputs, labels)

                val_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                val_acc += (predicted == labels).sum().item()

        # calculate average loss and accuracy for training and validation sets
        train_loss /= len(trainloader_mnist)
        train_acc /= len(trainset_mnist)
        val_loss /= len(valloader_mnist)
        val_acc /= len(valset_mnist)
        train_losses[i].append(train_loss)
        train_accs[i].append(train_acc)
        val_losses[i].append(val_loss)
        val_accs[i].append(val_acc)

        # evaluate the model on test set
        model.eval()
        with torch.no_grad():
            test_acc = 0.0
            for j, data in enumerate(testloader_mnist, 0):
                inputs, labels = data
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                test_acc += (predicted == labels).sum().item()

            test_acc /= len(testset_mnist)
            test_accs[i].append(test_acc)

        # print the loss and accuracy for each epoch
        print(f'Model {i+1} | Epoch {epoch+1:3d} | Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f} | Test Acc: {test_acc:.4f}')


# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(train_losses[i], color=colors[i], label=f'Train Loss ({models[i].__class__.__name__})')
    plt.plot(val_losses[i], '--', color=colors[i], label=f'Val Loss ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('Training and Validation Loss over Epochs')
plt.savefig("Q21figure1.pdf")
plt.show()


# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(train_accs[i], color=colors[i], label=f'Train Acc ({models[i].__class__.__name__})')
    plt.plot(val_accs[i], '--', color=colors[i], label=f'Val Acc ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.title('Training and Validation Accuracy over Epochs')
plt.savefig("Q21figure2.pdf")
plt.show()

# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(test_accs[i], color=colors[i], label=f'Test Acc ({models[i].__class__.__name__})')
    # plt.plot(val_accs[i], '--', color=colors[i], label=f'Val Loss ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('Test Accuracy over Epochs')
plt.savefig("Q21figure3.pdf")
plt.show()

# evaluate the model on test set
for i, model in enumerate(models):
    print(f"Test model {i+1}")
    model.eval()
    with torch.no_grad():
      test_acc = 0.0
      all_targets = []
      all_predictions = []
      for i, data in enumerate(testloader_mnist, 0):
          inputs, labels = data
          outputs = model(inputs)
          _, predicted = torch.max(outputs.data, 1)
          test_acc += (predicted == labels).sum().item()
          all_targets.extend(labels.cpu().numpy())
          all_predictions.extend(predicted.cpu().numpy())

      # calculate accuracy on test set
      test_acc /= len(testset_mnist)
      test_accs.append(test_acc)
      
      # print classification report
      print('Classification report:')
      print(classification_report(all_targets, all_predictions))
      
      # plot confusion matrix
      cm = confusion_matrix(all_targets, all_predictions)
      ax = sns.heatmap(cm, annot=True, fmt='g', cmap='Blues', cbar=False)
      ax.set_xlabel('Predicted labels')
      ax.set_ylabel('True labels')
      ax.xaxis.set_ticklabels(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])
      ax.yaxis.set_ticklabels(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])
      plt.title(f'Confusion Matrix ({model.__class__.__name__})')
      plt.savefig(f"confusion_matrix_{model.__class__.__name__}.pdf")
      plt.show()

"""### Test2"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# Load the preprocessed datasets
trainset_mnist = torch.load('./data/preprocessed/trainset_cifar10.pt')
valset_mnist = torch.load('./data/preprocessed/valset_cifar10.pt')
testset_mnist = torch.load('./data/preprocessed/testset_cifar10.pt')

# Define hyperparameters
batch_size = 128
learning_rate = 0.02
momentum = 0.9
weight_decay = 0.000005
num_epochs = 150

# Create data loaders
trainloader_mnist = DataLoader(trainset_mnist, batch_size=batch_size, shuffle=True)
valloader_mnist = DataLoader(valset_mnist, batch_size=batch_size, shuffle=False)
testloader_mnist = DataLoader(testset_mnist, batch_size=batch_size, shuffle=False)

# Define the model
class SCNNB(nn.Module):
    def __init__(self):
        super(SCNNB, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-a
class SCNNB_a(nn.Module):
    def __init__(self):
        super(SCNNB_a, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-b
class SCNNB_b(nn.Module):
    def __init__(self):
        super(SCNNB_b, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# define the models
models = [SCNNB(), SCNNB_a(), SCNNB_b()]
colors = ['blue', 'green', 'red']

# define the training loop
train_losses = [[] for _ in range(len(models))]
train_accs = [[] for _ in range(len(models))]
val_losses = [[] for _ in range(len(models))]
val_accs = [[] for _ in range(len(models))]
test_accs = [[] for _ in range(len(models))]

for i, model in enumerate(models):
    print(f"Training model {i+1}")
    for epoch in range(num_epochs):

        criterion = nn.CrossEntropyLoss()
        optimizer = optim.SGD(model.parameters(), lr=0.02, momentum=0.9, weight_decay=0.000005)
        train_loss = 0.0
        train_total = 0
        train_correct = 0
        val_loss = 0.0
        val_total = 0
        val_correct = 0
        train_acc = 0.0
        val_acc = 0.0

        # train the model on training set
        model.train()
        for j, data in enumerate(trainloader_mnist, 0):
            inputs, labels = data
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            train_acc += (predicted == labels).sum().item()

        # evaluate the model on validation set
        model.eval()
        with torch.no_grad():
            for j, data in enumerate(valloader_mnist, 0):
                inputs, labels = data
                outputs = model(inputs)
                loss = criterion(outputs, labels)

                val_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                val_acc += (predicted == labels).sum().item()

        # calculate average loss and accuracy for training and validation sets
        train_loss /= len(trainloader_mnist)
        train_acc /= len(trainset_mnist)
        val_loss /= len(valloader_mnist)
        val_acc /= len(valset_mnist)
        train_losses[i].append(train_loss)
        train_accs[i].append(train_acc)
        val_losses[i].append(val_loss)
        val_accs[i].append(val_acc)

        # evaluate the model on test set
        model.eval()
        with torch.no_grad():
            test_acc = 0.0
            for j, data in enumerate(testloader_mnist, 0):
                inputs, labels = data
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                test_acc += (predicted == labels).sum().item()

            test_acc /= len(testset_mnist)
            test_accs[i].append(test_acc)

        # print the loss and accuracy for each epoch
        print(f'Model {i+1} | Epoch {epoch+1:3d} | Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f} | Test Acc: {test_acc:.4f}')


# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(train_losses[i], color=colors[i], label=f'Train Loss ({models[i].__class__.__name__})')
    plt.plot(val_losses[i], '--', color=colors[i], label=f'Val Loss ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('Training and Validation Loss over Epochs')
plt.savefig("Q21figure1.pdf")
plt.show()


# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(train_accs[i], color=colors[i], label=f'Train Acc ({models[i].__class__.__name__})')
    plt.plot(val_accs[i], '--', color=colors[i], label=f'Val Acc ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.title('Training and Validation Loss over Epochs')
plt.savefig("Q21figure2.pdf")
plt.show()

# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(test_accs[i], color=colors[i], label=f'Test Acc ({models[i].__class__.__name__})')
    # plt.plot(val_accs[i], '--', color=colors[i], label=f'Val Loss ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('Test Accuracy over Epochs')
plt.savefig("Q21figure3.pdf")
plt.show()

# evaluate the model on test set
for i, model in enumerate(models):
    print(f"Test model {i+1}")
    model.eval()
    with torch.no_grad():
      test_acc = 0.0
      all_targets = []
      all_predictions = []
      for i, data in enumerate(testloader_mnist, 0):
          inputs, labels = data
          outputs = model(inputs)
          _, predicted = torch.max(outputs.data, 1)
          test_acc += (predicted == labels).sum().item()
          all_targets.extend(labels.cpu().numpy())
          all_predictions.extend(predicted.cpu().numpy())

      # calculate accuracy on test set
      test_acc /= len(testset_mnist)
      test_accs.append(test_acc)
      
      # print classification report
      print('Classification report:')
      print(classification_report(all_targets, all_predictions))
      
      # plot confusion matrix
      for i in range(len(models)):
        cm = confusion_matrix(all_targets, all_predictions)
        ax = sns.heatmap(cm, annot=True, fmt='g', cmap='Blues', cbar=False)
        ax.set_xlabel('Predicted labels')
        ax.set_ylabel('True labels')
        ax.xaxis.set_ticklabels(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])
        ax.yaxis.set_ticklabels(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])
        plt.title(f'Confusion Matrix ({models[i].__class__.__name__})')
        plt.savefig(f"Q21figure4_{models[i].__class__.__name__}.pdf")
        plt.show()

"""# OTHER(!)"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# Load the preprocessed datasets
trainset_mnist = torch.load('./data/preprocessed/trainset_fashion_mnist.pt')
valset_mnist = torch.load('./data/preprocessed/valset_fashion_mnist.pt')
testset_mnist = torch.load('./data/preprocessed/testset_fashion_mnist.pt')

# Define hyperparameters
batch_size = 128
learning_rate = 0.02
momentum = 0.9
weight_decay = 0.000005
num_epochs = 2

# Create data loaders
trainloader_mnist = DataLoader(trainset_mnist, batch_size=batch_size, shuffle=True)
valloader_mnist = DataLoader(valset_mnist, batch_size=batch_size, shuffle=False)
testloader_mnist = DataLoader(testset_mnist, batch_size=batch_size, shuffle=False)

# Define the model
class SCNNB(nn.Module):
    def __init__(self):
        super(SCNNB, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-a
class SCNNB_a(nn.Module):
    def __init__(self):
        super(SCNNB_a, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-b
class SCNNB_b(nn.Module):
    def __init__(self):
        super(SCNNB_b, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# define the models
models = [SCNNB(), SCNNB_a(), SCNNB_b()]
colors = ['blue', 'green', 'red']

# define the optimizer and loss function
import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=weight_decay)

# # define the optimizer and loss function
# import torch.optim as optim
# criterion = nn.CrossEntropyLoss()
# optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=weight_decay)

# # define the training loop
train_losses = [[] for _ in range(len(models))]
train_accs = [[] for _ in range(len(models))]
val_losses = [[] for _ in range(len(models))]
val_accs = [[] for _ in range(len(models))]
test_accs = [[] for _ in range(len(models))]

for i, model in enumerate(models):
    print(f"Training model {i+1}")
    for epoch in range(num_epochs):
        train_loss = 0.0
        train_total = 0
        train_correct = 0
        val_loss = 0.0
        val_total = 0
        val_correct = 0
        train_acc = 0.0
        val_acc = 0.0

        # train the model on training set
        model.train()
        for j, data in enumerate(trainloader_mnist, 0):
            inputs, labels = data
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            train_acc += (predicted == labels).sum().item()

        # evaluate the model on validation set
        model.eval()
        with torch.no_grad():
            for j, data in enumerate(valloader_mnist, 0):
                inputs, labels = data
                outputs = model(inputs)
                loss = criterion(outputs, labels)

                val_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                val_acc += (predicted == labels).sum().item()

        # calculate average loss and accuracy for training and validation sets
        train_loss /= len(trainloader_mnist)
        train_acc /= len(trainset_mnist)
        val_loss /= len(valloader_mnist)
        val_acc /= len(valset_mnist)
        train_losses[i].append(train_loss)
        train_accs[i].append(train_acc)
        val_losses[i].append(val_loss)
        val_accs[i].append(val_acc)

        # evaluate the model on test set
        model.eval()
        with torch.no_grad():
            test_acc = 0.0
            for j, data in enumerate(testloader_mnist, 0):
                inputs, labels = data
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                test_acc += (predicted == labels).sum().item()

            test_acc /= len(testset_mnist)
            test_accs[i].append(test_acc)

        # print the loss and accuracy for each epoch
        print(f'Model {i+1} | Epoch {epoch+1:3d} | Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f} | Test Acc: {test_acc:.4f}')


# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(train_losses[i], color=colors[i], label=f'Train Loss ({models[i].__class__.__name__})')
    plt.plot(val_losses[i], '--', color=colors[i], label=f'Val Loss ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('Training and Validation Loss over Epochs')
plt.savefig("Q21figure1.pdf")
plt.show()


# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(train_accs[i], color=colors[i], label=f'Train Acc ({models[i].__class__.__name__})')
    plt.plot(val_accs[i], '--', color=colors[i], label=f'Val Acc ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.title('Training and Validation Loss over Epochs')
plt.savefig("Q21figure1.pdf")
plt.show()

# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(test_accs[i], color=colors[i], label=f'Test Acc ({models[i].__class__.__name__})')
    # plt.plot(val_accs[i], '--', color=colors[i], label=f'Val Loss ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('Test Accuracy over Epochs')
plt.savefig("Q21figure3.pdf")
plt.show()

# # plot the loss and accuracy for training and validation sets
# import matplotlib.pyplot as plt
# plt.plot(train_losses, label='Train Loss')
# plt.plot(val_losses, label='Val Loss')
# plt.xlabel('Epoch')
# plt.ylabel('Loss')
# plt.legend()
# plt.title('Training and Validation Loss over Epochs (SCNNB)')
# plt.savefig("Q21figure1.pdf")
# plt.show()

# plt.plot(train_accs, label='Train Acc')
# plt.plot(val_accs, label='Val Acc')
# plt.xlabel('Epoch')
# plt.ylabel('Accuracy')
# plt.legend()
# plt.title('Training and Validation Accuracy over Epochs (SCNNB)')
# plt.savefig("Q21figure2.pdf")
# plt.show()

# # plot the accuracy for test set
# plt.plot(test_accs, label='Test Acc')
# plt.xlabel('Epoch')
# plt.ylabel('Accuracy')
# plt.legend()
# plt.title('Test Accuracy Epochs (SCNNB)')
# plt.savefig("Q21figure3.pdf")
# plt.show()

# evaluate the model on test set
for i, model in enumerate(models):
    print(f"Training model {i+1}")
    model.eval()
    with torch.no_grad():
      test_acc = 0.0
      all_targets = []
      all_predictions = []
      for i, data in enumerate(testloader_mnist, 0):
          inputs, labels = data
          outputs = model(inputs)
          _, predicted = torch.max(outputs.data, 1)
          test_acc += (predicted == labels).sum().item()
          all_targets.extend(labels.cpu().numpy())
          all_predictions.extend(predicted.cpu().numpy())

      # calculate accuracy on test set
      test_acc /= len(testset_mnist)
      test_accs.append(test_acc)
      
      # print classification report
      print('Classification report:')
      print(classification_report(all_targets, all_predictions))
      
      # plot confusion matrix
      cm = confusion_matrix(all_targets, all_predictions)
      ax = sns.heatmap(cm, annot=True, fmt='g', cmap='Blues', cbar=False)
      ax.set_xlabel('Predicted labels')
      ax.set_ylabel('True labels')
      ax.xaxis.set_ticklabels(['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot'])
      ax.yaxis.set_ticklabels(['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot'])
      plt.title('Confusion Matrix (SCNNB)')
      plt.savefig("Q21figure4.pdf")
      plt.show()

for i in range(len(models)):
        cm = confusion_matrix(all_targets, all_predictions)
        ax = sns.heatmap(cm, annot=True, fmt='g', cmap='Blues', cbar=False)
        ax.set_xlabel('Predicted labels')
        ax.set_ylabel('True labels')
        ax.xaxis.set_ticklabels(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])
        ax.yaxis.set_ticklabels(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])
        plt.title(f'Confusion Matrix ({models[i].__class__.__name__})')
        plt.savefig(f"Q21figure4_{models[i].__class__.__name__}.pdf")
        plt.show()

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# Load the preprocessed datasets
trainset_mnist = torch.load('./data/preprocessed/trainset_fashion_mnist.pt')
valset_mnist = torch.load('./data/preprocessed/valset_fashion_mnist.pt')
testset_mnist = torch.load('./data/preprocessed/testset_fashion_mnist.pt')

# Define hyperparameters
batch_size = 128
learning_rate = 0.02
momentum = 0.9
weight_decay = 0.000005
num_epochs = 2

# Create data loaders
trainloader_mnist = DataLoader(trainset_mnist, batch_size=batch_size, shuffle=True)
valloader_mnist = DataLoader(valset_mnist, batch_size=batch_size, shuffle=False)
testloader_mnist = DataLoader(testset_mnist, batch_size=batch_size, shuffle=False)

# Define the model
class SCNNB(nn.Module):
    def __init__(self):
        super(SCNNB, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-a
class SCNNB_a(nn.Module):
    def __init__(self):
        super(SCNNB_a, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-b
class SCNNB_b(nn.Module):
    def __init__(self):
        super(SCNNB_b, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# define the models
models = [SCNNB(), SCNNB_a(), SCNNB_b()]
colors = ['blue', 'green', 'red']

# define the optimizer and loss function
import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=weight_decay)

# define the training loop
train_losses = [[] for _ in range(len(models))]
train_accs = [[] for _ in range(len(models))]
val_losses = [[] for _ in range(len(models))]
val_accs = [[] for _ in range(len(models))]
test_accs = [[] for _ in range(len(models))]

for i, model in enumerate(models):
    print(f"Training model {i+1}")
    for epoch in range(num_epochs):
        train_loss = 0.0
        train_total = 0
        train_correct = 0
        val_loss = 0.0
        val_total = 0
        val_correct = 0
        train_loss = 0.0
        train_acc = 0.0
        val_loss = 0.0
        val_acc = 0.0

    # define the optimizer and loss function
    import torch.optim as optim
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.02, momentum=0.9, weight_decay=0.000005)
    train_losses = [[] for _ in range(len(models))]
    train_accs = [[] for _ in range(len(models))]
    val_losses = [[] for _ in range(len(models))]
    val_accs = [[] for _ in range(len(models))]
    test_accs = [[] for _ in range(len(models))]

    for epoch in range(num_epochs):
        train_loss = 0.0
        train_total = 0
        train_correct = 0
        val_loss = 0.0
        val_total = 0
        val_correct = 0
        train_loss = 0.0
        train_acc = 0.0
        val_loss = 0.0
        val_acc = 0.0

        # train the model on training set
        model.train()
        for j, data in enumerate(trainloader_mnist, 0):
            inputs, labels = data
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            train_acc += (predicted == labels).sum().item()

        # evaluate the model on validation set
        model.eval()
        with torch.no_grad():
            for j, data in enumerate(valloader_mnist, 0):
                inputs, labels = data
                outputs = model(inputs)
                loss = criterion(outputs, labels)

                val_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                val_acc += (predicted == labels).sum().item()

        # calculate average loss and accuracy for training and validation sets
        train_loss /= len(trainloader_mnist)
        train_acc /= len(trainset_mnist)
        val_loss /= len(valloader_mnist)
        val_acc /= len(valset_mnist)
        train_losses[i].append(train_loss)
        train_accs[i].append(train_acc)
        val_losses[i].append(val_loss)
        val_accs[i].append(val_acc)

        # evaluate the model on test set
        model.eval()
        with torch.no_grad():
            test_acc = 0.0
            for j, data in enumerate(testloader_mnist, 0):
                inputs, labels = data
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                test_acc += (predicted == labels).sum().item()

            test_acc /= len(testset_mnist)
            test_accs[i].append(test_acc)

        # print the loss and accuracy for each epoch
        print(f'Model {i+1} | Epoch {epoch+1:3d} | Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f} | Test Acc: {test_acc:.4f}')


# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(train_losses[i], color=colors[i], label=f'Train Loss ({models[i].__class__.__name__})')
    plt.plot(val_losses[i], '--', color=colors[i], label=f'Val Loss ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('Training and Validation Loss over Epochs')
plt.savefig("Q21figure1.pdf")
plt.show()


# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(train_accs[i], color=colors[i], label=f'Train Acc ({models[i].__class__.__name__})')
    plt.plot(val_accs[i], '--', color=colors[i], label=f'Val Acc ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.title('Training and Validation Loss over Epochs')
plt.savefig("Q21figure1.pdf")
plt.show()

# Plot the results for all models in the same graph
plt.figure(figsize=(8, 6))
for i in range(len(models)):
    plt.plot(test_accs[i], color=colors[i], label=f'Test Acc ({models[i].__class__.__name__})')
    # plt.plot(val_accs[i], '--', color=colors[i], label=f'Val Loss ({models[i].__class__.__name__})')

plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('Test Accuracy over Epochs')
plt.savefig("Q21figure3.pdf")
plt.show()

# # plot the loss and accuracy for training and validation sets
# import matplotlib.pyplot as plt
# plt.plot(train_losses, label='Train Loss')
# plt.plot(val_losses, label='Val Loss')
# plt.xlabel('Epoch')
# plt.ylabel('Loss')
# plt.legend()
# plt.title('Training and Validation Loss over Epochs (SCNNB)')
# plt.savefig("Q21figure1.pdf")
# plt.show()

# plt.plot(train_accs, label='Train Acc')
# plt.plot(val_accs, label='Val Acc')
# plt.xlabel('Epoch')
# plt.ylabel('Accuracy')
# plt.legend()
# plt.title('Training and Validation Accuracy over Epochs (SCNNB)')
# plt.savefig("Q21figure2.pdf")
# plt.show()

# # plot the accuracy for test set
# plt.plot(test_accs, label='Test Acc')
# plt.xlabel('Epoch')
# plt.ylabel('Accuracy')
# plt.legend()
# plt.title('Test Accuracy Epochs (SCNNB)')
# plt.savefig("Q21figure3.pdf")
# plt.show()

# evaluate the model on test set
for i, model in enumerate(models):
    print(f"Test model {i+1}")
    model.eval()
    with torch.no_grad():
      test_acc = 0.0
      all_targets = []
      all_predictions = []
      for i, data in enumerate(testloader_mnist, 0):
          inputs, labels = data
          outputs = model(inputs)
          _, predicted = torch.max(outputs.data, 1)
          test_acc += (predicted == labels).sum().item()
          all_targets.extend(labels.cpu().numpy())
          all_predictions.extend(predicted.cpu().numpy())

      # calculate accuracy on test set
      test_acc /= len(testset_mnist)
      test_accs.append(test_acc)
      
      # print classification report
      print('Classification report:')
      print(classification_report(all_targets, all_predictions))
      
      # plot confusion matrix
      cm = confusion_matrix(all_targets, all_predictions)
      ax = sns.heatmap(cm, annot=True, fmt='g', cmap='Blues', cbar=False)
      ax.set_xlabel('Predicted labels')
      ax.set_ylabel('True labels')
      ax.xaxis.set_ticklabels(['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot'])
      ax.yaxis.set_ticklabels(['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot'])
      plt.title('Confusion Matrix (SCNNB)')
      plt.savefig("Q21figure4.pdf")
      plt.show()

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# Load the preprocessed datasets
trainset_mnist = torch.load('./data/preprocessed/trainset_fashion_mnist.pt')
valset_mnist = torch.load('./data/preprocessed/valset_fashion_mnist.pt')
testset_mnist = torch.load('./data/preprocessed/testset_fashion_mnist.pt')

# Define hyperparameters
batch_size = 128
learning_rate = 0.02
momentum = 0.9
weight_decay = 0.000005
num_epochs = 300

# Create data loaders
trainloader_mnist = DataLoader(trainset_mnist, batch_size=batch_size, shuffle=True)
valloader_mnist = DataLoader(valset_mnist, batch_size=batch_size, shuffle=False)
testloader_mnist = DataLoader(testset_mnist, batch_size=batch_size, shuffle=False)

# Define the model
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-a
class SCNNB_a(nn.Module):
    def __init__(self):
        super(SCNNB_a, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-b
class SCNNB_b(nn.Module):
    def __init__(self):
        super(SCNNB_b, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x


# Initialize the model and the optimizer
model = MyModel()
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=weight_decay)

# Define the loss function
criterion = nn.CrossEntropyLoss()

# Define lists to store the loss and accuracy for the train and validation sets
train_losses = []
train_accs = []
val_losses = []
val_accs = []

# Train the model
for epoch in range(num_epochs):
    train_loss = 0.0
    train_total = 0
    train_correct = 0
    val_loss = 0.0
    val_total = 0
    val_correct = 0
    
# define the optimizer and loss function
import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.02, momentum=0.9, weight_decay=0.000005)

# define the training loop
# num_epochs = 100
train_losses = []
train_accs = []
val_losses = []
val_accs = []
test_accs = []

for epoch in range(num_epochs):
    train_loss = 0.0
    train_acc = 0.0
    val_loss = 0.0
    val_acc = 0.0

    # train the model on training set
    model.train()
    for i, data in enumerate(trainloader_mnist, 0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        train_acc += (predicted == labels).sum().item()

    # evaluate the model on validation set
    model.eval()
    with torch.no_grad():
        for i, data in enumerate(valloader_mnist, 0):
            inputs, labels = data
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            val_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            val_acc += (predicted == labels).sum().item()

    # calculate average loss and accuracy for training and validation sets
    train_loss /= len(trainloader_mnist)
    train_acc /= len(trainset_mnist)
    val_loss /= len(valloader_mnist)
    val_acc /= len(valset_mnist)
    train_losses.append(train_loss)
    train_accs.append(train_acc)
    val_losses.append(val_loss)
    val_accs.append(val_acc)

    # evaluate the model on test set
    model.eval()
    with torch.no_grad():
        test_acc = 0.0
        for i, data in enumerate(testloader_mnist, 0):
            inputs, labels = data
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            test_acc += (predicted == labels).sum().item()

        test_acc /= len(testset_mnist)
        test_accs.append(test_acc)

    # print the loss and accuracy for each epoch
    print(f'Epoch {epoch+1:3d} | Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f} | Test Acc: {test_acc:.4f}')

# plot the loss and accuracy for training and validation sets
import matplotlib.pyplot as plt
plt.plot(train_losses, label='Train Loss')
plt.plot(val_losses, label='Val Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('Training and Validation Loss over Epochs (SCNNB)')
plt.savefig("Q21figure1.pdf")
plt.show()

plt.plot(train_accs, label='Train Acc')
plt.plot(val_accs, label='Val Acc')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.title('Training and Validation Accuracy over Epochs (SCNNB)')
plt.savefig("Q21figure2.pdf")
plt.show()

# plot the accuracy for test set
plt.plot(test_accs, label='Test Acc')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.title('Test Accuracy Epochs (SCNNB)')
plt.savefig("Q21figure3.pdf")
plt.show()

# # evaluate the model on test set
# model.eval()
# with torch.no_grad():
#     all_labels = []
#     all_predictions = []
#     for i, data in enumerate(testloader, 0):
#         inputs, labels = data
#         outputs = model(inputs)
#         _, predicted = torch.max(outputs.data, 1)
#         all_labels.extend

# evaluate the model on test set
model.eval()
with torch.no_grad():
    test_acc = 0.0
    all_targets = []
    all_predictions = []
    for i, data in enumerate(testloader_mnist, 0):
        inputs, labels = data
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        test_acc += (predicted == labels).sum().item()
        all_targets.extend(labels.cpu().numpy())
        all_predictions.extend(predicted.cpu().numpy())

    # calculate accuracy on test set
    test_acc /= len(testset_mnist)
    test_accs.append(test_acc)
    
    # print classification report
    print('Classification report:')
    print(classification_report(all_targets, all_predictions))
    
    from sklearn.metrics import confusion_matrix
    import seaborn as sns
    import matplotlib.pyplot as plt
    import numpy as np

    # Define the true labels and predicted labels
    true_labels = np.array(test_labels) # assuming `test_labels` contains true labels
    predicted_labels = np.array(test_predictions) # assuming `test_predictions` contains predicted labels

    # Compute confusion matrix
    cm = confusion_matrix(true_labels, predicted_labels)

    # Define class names
    class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

    # Plot confusion matrix
    fig, ax = plt.subplots(figsize=(8, 8))
    sns.heatmap(cm, annot=True, cmap='Blues', cbar=False, xticklabels=class_names, yticklabels=class_names, ax=ax)
    ax.set_xlabel('Predicted labels')
    ax.set_ylabel('True labels')
    ax.set_title('Confusion Matrix (SCNNB)')

    # Save and show the plot
    plt.savefig("Q21figure4.pdf")
    plt.show()

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# Load the preprocessed datasets
trainset_mnist = torch.load('./data/preprocessed/trainset_fashion_mnist.pt')
valset_mnist = torch.load('./data/preprocessed/valset_fashion_mnist.pt')
testset_mnist = torch.load('./data/preprocessed/testset_fashion_mnist.pt')

# Define hyperparameters
batch_size = 128
learning_rate = 0.02
momentum = 0.9
weight_decay = 0.000005
num_epochs = 300

# Create data loaders
trainloader_mnist = DataLoader(trainset_mnist, batch_size=batch_size, shuffle=True)
valloader_mnist = DataLoader(valset_mnist, batch_size=batch_size, shuffle=False)
testloader_mnist = DataLoader(testset_mnist, batch_size=batch_size, shuffle=False)

# Define the model
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-a
class SCNNB_a(nn.Module):
    def __init__(self):
        super(SCNNB_a, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-b
class SCNNB_b(nn.Module):
    def __init__(self):
        super(SCNNB_b, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x


# Initialize the model and the optimizer
model = SCNNB_a()
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=weight_decay)

# Define the loss function
criterion = nn.CrossEntropyLoss()

# Define lists to store the loss and accuracy for the train and validation sets
train_losses = []
train_accs = []
val_losses = []
val_accs = []

# Train the model
for epoch in range(num_epochs):
    train_loss = 0.0
    train_total = 0
    train_correct = 0
    val_loss = 0.0
    val_total = 0
    val_correct = 0
    
# define the optimizer and loss function
import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.02, momentum=0.9, weight_decay=0.000005)

# define the training loop
# num_epochs = 100
train_losses = []
train_accs = []
val_losses = []
val_accs = []
test_accs = []

for epoch in range(num_epochs):
    train_loss = 0.0
    train_acc = 0.0
    val_loss = 0.0
    val_acc = 0.0

    # train the model on training set
    model.train()
    for i, data in enumerate(trainloader_mnist, 0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        train_acc += (predicted == labels).sum().item()

    # evaluate the model on validation set
    model.eval()
    with torch.no_grad():
        for i, data in enumerate(valloader_mnist, 0):
            inputs, labels = data
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            val_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            val_acc += (predicted == labels).sum().item()

    # calculate average loss and accuracy for training and validation sets
    train_loss /= len(trainloader_mnist)
    train_acc /= len(trainset_mnist)
    val_loss /= len(valloader_mnist)
    val_acc /= len(valset_mnist)
    train_losses.append(train_loss)
    train_accs.append(train_acc)
    val_losses.append(val_loss)
    val_accs.append(val_acc)

    # evaluate the model on test set
    model.eval()
    with torch.no_grad():
        test_acc = 0.0
        for i, data in enumerate(testloader_mnist, 0):
            inputs, labels = data
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            test_acc += (predicted == labels).sum().item()

        test_acc /= len(testset_mnist)
        test_accs.append(test_acc)

    # print the loss and accuracy for each epoch
    print(f'Epoch {epoch+1:3d} | Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f} | Test Acc: {test_acc:.4f}')

# plot the loss and accuracy for training and validation sets
import matplotlib.pyplot as plt
plt.plot(train_losses, label='Train Loss')
plt.plot(val_losses, label='Val Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('Training and Validation Loss over Epochs (SCNNB-a)')
plt.savefig("Q21figure1111.pdf")
plt.show()

plt.plot(train_accs, label='Train Acc')
plt.plot(val_accs, label='Val Acc')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.title('Training and Validation Accuracy over Epochs (SCNNB-a)')
plt.savefig("Q21figure2111.pdf")
plt.show()

# plot the accuracy for test set
plt.plot(test_accs, label='Test Acc')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.title('Test Accuracy Epochs (SCNNB-a)')
plt.savefig("Q21figure3111.pdf")
plt.show()

# # evaluate the model on test set
# model.eval()
# with torch.no_grad():
#     all_labels = []
#     all_predictions = []
#     for i, data in enumerate(testloader, 0):
#         inputs, labels = data
#         outputs = model(inputs)
#         _, predicted = torch.max(outputs.data, 1)
#         all_labels.extend

# evaluate the model on test set
model.eval()
with torch.no_grad():
    test_acc = 0.0
    all_targets = []
    all_predictions = []
    for i, data in enumerate(testloader_mnist, 0):
        inputs, labels = data
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        test_acc += (predicted == labels).sum().item()
        all_targets.extend(labels.cpu().numpy())
        all_predictions.extend(predicted.cpu().numpy())

    # calculate accuracy on test set
    test_acc /= len(testset_mnist)
    test_accs.append(test_acc)
    
    # print classification report
    print('Classification report:')
    print(classification_report(all_targets, all_predictions))
    
    from sklearn.metrics import confusion_matrix
    import seaborn as sns
    import matplotlib.pyplot as plt
    import numpy as np

    # Define the true labels and predicted labels
    true_labels = np.array(test_labels) # assuming `test_labels` contains true labels
    predicted_labels = np.array(test_predictions) # assuming `test_predictions` contains predicted labels

    # Compute confusion matrix
    cm = confusion_matrix(true_labels, predicted_labels)

    # Define class names
    class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

    # Plot confusion matrix
    fig, ax = plt.subplots(figsize=(8, 8))
    sns.heatmap(cm, annot=True, cmap='Blues', cbar=False, xticklabels=class_names, yticklabels=class_names, ax=ax)
    ax.set_xlabel('Predicted labels')
    ax.set_ylabel('True labels')
    ax.set_title('Confusion Matrix (SCNNB-a)')

    # Save and show the plot
    plt.savefig("Q21figure4111.pdf")
    plt.show()

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# Load the preprocessed datasets
trainset_mnist = torch.load('./data/preprocessed/trainset_fashion_mnist.pt')
valset_mnist = torch.load('./data/preprocessed/valset_fashion_mnist.pt')
testset_mnist = torch.load('./data/preprocessed/testset_fashion_mnist.pt')

# Define hyperparameters
batch_size = 128
learning_rate = 0.02
momentum = 0.9
weight_decay = 0.000005
num_epochs = 300

# Create data loaders
trainloader_mnist = DataLoader(trainset_mnist, batch_size=batch_size, shuffle=True)
valloader_mnist = DataLoader(valset_mnist, batch_size=batch_size, shuffle=False)
testloader_mnist = DataLoader(testset_mnist, batch_size=batch_size, shuffle=False)

# Define the model
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-a
class SCNNB_a(nn.Module):
    def __init__(self):
        super(SCNNB_a, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-b
class SCNNB_b(nn.Module):
    def __init__(self):
        super(SCNNB_b, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x


# Initialize the model and the optimizer
model = SCNNB_b()
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=weight_decay)

# Define the loss function
criterion = nn.CrossEntropyLoss()

# Define lists to store the loss and accuracy for the train and validation sets
train_losses = []
train_accs = []
val_losses = []
val_accs = []

# Train the model
for epoch in range(num_epochs):
    train_loss = 0.0
    train_total = 0
    train_correct = 0
    val_loss = 0.0
    val_total = 0
    val_correct = 0
    
# define the optimizer and loss function
import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.02, momentum=0.9, weight_decay=0.000005)

# define the training loop
# num_epochs = 100
train_losses = []
train_accs = []
val_losses = []
val_accs = []
test_accs = []

for epoch in range(num_epochs):
    train_loss = 0.0
    train_acc = 0.0
    val_loss = 0.0
    val_acc = 0.0

    # train the model on training set
    model.train()
    for i, data in enumerate(trainloader_mnist, 0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        train_acc += (predicted == labels).sum().item()

    # evaluate the model on validation set
    model.eval()
    with torch.no_grad():
        for i, data in enumerate(valloader_mnist, 0):
            inputs, labels = data
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            val_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            val_acc += (predicted == labels).sum().item()

    # calculate average loss and accuracy for training and validation sets
    train_loss /= len(trainloader_mnist)
    train_acc /= len(trainset_mnist)
    val_loss /= len(valloader_mnist)
    val_acc /= len(valset_mnist)
    train_losses.append(train_loss)
    train_accs.append(train_acc)
    val_losses.append(val_loss)
    val_accs.append(val_acc)

    # evaluate the model on test set
    model.eval()
    with torch.no_grad():
        test_acc = 0.0
        for i, data in enumerate(testloader_mnist, 0):
            inputs, labels = data
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            test_acc += (predicted == labels).sum().item()

        test_acc /= len(testset_mnist)
        test_accs.append(test_acc)

    # print the loss and accuracy for each epoch
    print(f'Epoch {epoch+1:3d} | Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f} | Test Acc: {test_acc:.4f}')

# plot the loss and accuracy for training and validation sets
import matplotlib.pyplot as plt
plt.plot(train_losses, label='Train Loss')
plt.plot(val_losses, label='Val Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.title('Training and Validation Loss over Epochs (SCNNB-b)')
plt.savefig("Q21figure11222211.pdf")
plt.show()

plt.plot(train_accs, label='Train Acc')
plt.plot(val_accs, label='Val Acc')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.title('Training and Validation Accuracy over Epochs (SCNNB-b)')
plt.savefig("Q21figure2122211.pdf")
plt.show()

# plot the accuracy for test set
plt.plot(test_accs, label='Test Acc')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.title('Test Accuracy Epochs (SCNNB-b)')
plt.savefig("Q21figure3122211.pdf")
plt.show()

# # evaluate the model on test set
# model.eval()
# with torch.no_grad():
#     all_labels = []
#     all_predictions = []
#     for i, data in enumerate(testloader, 0):
#         inputs, labels = data
#         outputs = model(inputs)
#         _, predicted = torch.max(outputs.data, 1)
#         all_labels.extend

# evaluate the model on test set
model.eval()
with torch.no_grad():
    test_acc = 0.0
    all_targets = []
    all_predictions = []
    for i, data in enumerate(testloader_mnist, 0):
        inputs, labels = data
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        test_acc += (predicted == labels).sum().item()
        all_targets.extend(labels.cpu().numpy())
        all_predictions.extend(predicted.cpu().numpy())

    # calculate accuracy on test set
    test_acc /= len(testset_mnist)
    test_accs.append(test_acc)
    
    # print classification report
    print('Classification report:')
    print(classification_report(all_targets, all_predictions))
    
    from sklearn.metrics import confusion_matrix
    import seaborn as sns
    import matplotlib.pyplot as plt
    import numpy as np

    # Define the true labels and predicted labels
    true_labels = np.array(test_labels) # assuming `test_labels` contains true labels
    predicted_labels = np.array(test_predictions) # assuming `test_predictions` contains predicted labels

    # Compute confusion matrix
    cm = confusion_matrix(true_labels, predicted_labels)

    # Define class names
    class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

    # Plot confusion matrix
    fig, ax = plt.subplots(figsize=(8, 8))
    sns.heatmap(cm, annot=True, cmap='Blues', cbar=False, xticklabels=class_names, yticklabels=class_names, ax=ax)
    ax.set_xlabel('Predicted labels')
    ax.set_ylabel('True labels')
    ax.set_title('Confusion Matrix (SCNNB-b)')

    # Save and show the plot
    plt.savefig("Q21figure4112221.pdf")
    plt.show()

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import classification_report, confusion_matrix

# Load datasets
trainset = torch.load('./data/preprocessed/trainset_mnist.pt')
valset = torch.load('./data/preprocessed/valset_mnist.pt')
testset = torch.load('./data/preprocessed/testset_mnist.pt')

trainloader = DataLoader(trainset, batch_size=128, shuffle=True)
valloader = DataLoader(valset, batch_size=128, shuffle=False)
testloader = DataLoader(testset, batch_size=128, shuffle=False)

# Define the model
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.bn2 = nn.BatchNorm2d(64)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 5 * 5, 1280)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(1280, 10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.pool1(self.relu(self.bn1(self.conv1(x))))
        x = self.pool2(self.relu(self.bn2(self.conv2(x))))
        x = x.view(-1, 64 * 5 * 5)
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.softmax(self.fc2(x))
        return x

model = Net()

# Set up the loss function, optimizer, and hyperparameters
criterion = nn.CrossLoss()
optimizer = optim.SGD(model.parameters(), lr=0.02, momentum=0.9, weight_decay=0.000005)
epochs = 300

# Train the model
train_loss, val_loss, train_acc, val_acc = [], [], [], []

for epoch in range(epochs):
    # Training
    running_loss, correct, total = 0, 0, 0
    model.train()
    for inputs, labels in trainloader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()
    train_loss.append(running_loss / len(trainloader))
    train_acc.append(100 * correct / total)

    # Validation
    running_loss, correct, total = 0, 0, 0
    model.eval()
    with torch.no_grad():
        for inputs, labels in valloader:
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
    val_loss.append(running_loss / len(valloader))
    val_acc.append(100 * correct / total)

# Plot Loss and Accuracy for train and validation sets
plt.figure()
plt.plot(train_loss, label='Train Loss')
plt.plot(val_loss, label='Validation Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.show()

plt.figure()
plt.plot(train_acc, label='Train Accuracy')
plt.plot(val_acc, label='Validation Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend()
plt.show()

# Evaluate on test set
model.eval()
test_acc, y_true, y_pred = 0, [], []
with torch.no_grad():
    for inputs, labels in testloader:
        outputs = model(inputs)
        _, predicted = outputs.max(1)
        test_acc += predicted.eq(labels).sum().item()
        y_true.extend(labels.tolist())
        y_pred.extend(predicted.tolist())

test_acc = 100 * test_acc / len(testset)
print(f'Test Accuracy: {test_acc:.2f}%')

# Print classification report
print(classification_report(y_true, y_pred))


# plot confusion matrix
    cm = confusion_matrix(all_targets, all_predictions)
    ax = sns.heatmap(cm, annot=True, fmt='g', cmap='Blues', cbar=False)
    ax.set_xlabel('Predicted labels')
    ax.set_ylabel('True labels')
    ax.xaxis.set_ticklabels(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])
    ax.yaxis.set_ticklabels(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])
    plt.title('Confusion Matrix')
    plt.show()

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# Load the preprocessed datasets
trainset_mnist = torch.load('./data/preprocessed/trainset_mnist.pt')
valset_mnist = torch.load('./data/preprocessed/valset_mnist.pt')
testset_mnist = torch.load('./data/preprocessed/testset_mnist.pt')

# Define hyperparameters
batch_size = 128
learning_rate = 0.02
momentum = 0.9
weight_decay = 0.000005
num_epochs = 3

# Create data loaders
trainloader_mnist = DataLoader(trainset_mnist, batch_size=batch_size, shuffle=True)
valloader_mnist = DataLoader(valset_mnist, batch_size=batch_size, shuffle=False)
testloader_mnist = DataLoader(testset_mnist, batch_size=batch_size, shuffle=False)

# Define the model
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# Initialize the model and the optimizer
model = MyModel()
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=weight_decay)

# Define the loss function
criterion = nn.CrossEntropyLoss()

# Define lists to store the loss and accuracy for the train and validation sets
train_losses = []
train_accs = []
val_losses = []
val_accs = []

# Train the model
for epoch in range(num_epochs):
    train_loss = 0.0
    train_total = 0
    train_correct = 0
    val_loss = 0.0
    val_total = 0
    val_correct = 0
    
# define the optimizer and loss function
import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.02, momentum=0.9, weight_decay=0.000005)

# define the training loop
num_epochs = 3
train_losses = []
train_accs = []
val_losses = []
val_accs = []
test_accs = []

for epoch in range(num_epochs):
    train_loss = 0.0
    train_acc = 0.0
    val_loss = 0.0
    val_acc = 0.0

    # train the model on training set
    model.train()
    for i, data in enumerate(trainloader_mnist, 0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        train_acc += (predicted == labels).sum().item()

    # evaluate the model on validation set
    model.eval()
    with torch.no_grad():
        for i, data in enumerate(valloader_mnist, 0):
            inputs, labels = data
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            val_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            val_acc += (predicted == labels).sum().item()

    # calculate average loss and accuracy for training and validation sets
    train_loss /= len(trainloader_mnist)
    train_acc /= len(trainset_mnist)
    val_loss /= len(valloader_mnist)
    val_acc /= len(valset_mnist)
    train_losses.append(train_loss)
    train_accs.append(train_acc)
    val_losses.append(val_loss)
    val_accs.append(val_acc)

    # evaluate the model on test set
    model.eval()
    with torch.no_grad():
        test_acc = 0.0
        for i, data in enumerate(testloader_mnist, 0):
            inputs, labels = data
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            test_acc += (predicted == labels).sum().item()

        test_acc /= len(testset_mnist)
        test_accs.append(test_acc)

    # print the loss and accuracy for each epoch
    print(f'Epoch {epoch+1:3d} | Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f} | Test Acc: {test_acc:.4f}')

# plot the loss and accuracy for training and validation sets
import matplotlib.pyplot as plt
plt.plot(train_losses, label='Train Loss')
plt.plot(val_losses, label='Val Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.show()

plt.plot(train_accs, label='Train Acc')
plt.plot(val_accs, label='Val Acc')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.show()

# plot the accuracy for test set
plt.plot(test_accs, label='Test Acc')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.show()

# # evaluate the model on test set
# model.eval()
# with torch.no_grad():
#     all_labels = []
#     all_predictions = []
#     for i, data in enumerate(testloader, 0):
#         inputs, labels = data
#         outputs = model(inputs)
#         _, predicted = torch.max(outputs.data, 1)
#         all_labels.extend

# evaluate the model on test set
model.eval()
with torch.no_grad():
    test_acc = 0.0
    all_targets = []
    all_predictions = []
    for i, data in enumerate(testloader_mnist, 0):
        inputs, labels = data
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        test_acc += (predicted == labels).sum().item()
        all_targets.extend(labels.cpu().numpy())
        all_predictions.extend(predicted.cpu().numpy())

    # calculate accuracy on test set
    test_acc /= len(testset_mnist)
    test_accs.append(test_acc)
    
    # print classification report
    print('Classification report:')
    print(classification_report(all_targets, all_predictions))
    
    # plot confusion matrix
    cm = confusion_matrix(all_targets, all_predictions)
    ax = sns.heatmap(cm, annot=True, fmt='g', cmap='Blues', cbar=False)
    ax.set_xlabel('Predicted labels')
    ax.set_ylabel('True labels')
    ax.xaxis.set_ticklabels(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])
    ax.yaxis.set_ticklabels(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])
    plt.title('Confusion Matrix')
    plt.show()

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(64*5*5, 1280)
        self.relu3 = nn.ReLU()
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(1280, 10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

model = Model()

import torch.optim as optim
import numpy as np

# Define hyperparameters
learning_rate = 0.02
momentum = 0.9
weight_decay = 0.000005
num_epochs = 300
batch_size = 128

# Define loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=weight_decay)

# Define data loaders
train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True)
val_loader = torch.utils.data.DataLoader(valset, batch_size=batch_size, shuffle=False)
test_loader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False)

# Define arrays to store loss and accuracy values
train_loss = []
train_acc = []
val_loss = []
val_acc = []

# Train the model
for epoch in range(num_epochs):
    epoch_train_loss = 0
    epoch_train_acc = 0
    epoch_val_loss = 0
    epoch_val_acc = 0
    
    # Train for one epoch
    model.train()
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        epoch_train_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        epoch_train_acc += (predicted == labels).sum().item()
    epoch_train_loss /= len(train_loader)
    epoch_train_acc /= len(train_loader.dataset)
    
    # Evaluate on validation set
    model.eval()
    with torch.no_grad():
        for images, labels in val_loader:
            outputs = model(images)
            loss = criterion(outputs, labels)
            epoch_val_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            epoch_val_acc += (predicted == labels).sum().

import torch.nn as nn

class SCNNB(nn.Module):
    def __init__(self):
        super(SCNNB, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(64 * 5 * 5, 1280)
        self.relu3 = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(1280, 10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64 * 5 * 5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

model = SCNNB()

import torch.utils.data as data
import torch

# Define hyperparameters
learning_rate = 0.02
momentum = 0.9
weight_decay = 0.000005
num_epochs = 300
batch_size = 128

train_data_path = "/content/data/preprocessed/trainset_mnist.pt"
val_data_path = "/content/data/preprocessed/valset_mnist.pt"
test_data_path = "/content/data/preprocessed/testset_mnist.pt"

train_data = torch.load(train_data_path)
val_data = torch.load(val_data_path)
test_data = torch.load(test_data_path)

train_loader = data.DataLoader(train_data, batch_size=batch_size, shuffle=True)
val_loader = data.DataLoader(val_data, batch_size=batch_size, shuffle=False)
test_loader = data.DataLoader(test_data, batch_size=batch_size, shuffle=False)



# Define loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=weight_decay)

for epoch in range(num_epochs):
    # Training loop
    model.train()
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        if (i+1) % 100 == 0:
            print ('Epoch [{}/{}], Step [{}/{}], Train Loss: {:.4f}' 
                   .format(epoch+1, num_epochs, i+1, len(train_loader), loss.item()))

    # Validation loop
    model.eval()
    with torch.no_grad():
        val_loss = 0
        val_total = 0
        val_correct = 0
        for images, labels in val_loader:
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            val_total += labels.size(0)
            val_correct += (predicted == labels).sum().item()
            val_loss += criterion(outputs, labels).item() * labels.size(0)
        val_loss /= val_total
        val_accuracy = 100 * val_correct / val_total
        print ('Epoch [{}/{}], Validation Loss: {:.4f}, Validation Accuracy: {:.2f}%'.format(epoch+1, num_epochs, val_loss, val_accuracy))

# Test loop
model.eval()
with torch.no_grad():
    test_total = 0
    test_correct = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        test_total += labels.size(0)
        test_correct += (predicted == labels).sum().item()
    test_accuracy = 100 * test_correct / test_total
    print ('Test Accuracy: {:.2f}%'.format(test_accuracy))

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import matplotlib.pyplot as plt
from sklearn.metrics import classification_report, confusion_matrix
import seaborn as sns

# Define the model
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)
        self.bn1 = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.bn2 = nn.BatchNorm2d(64)
        self.fc1 = nn.Linear(64 * 5 * 5, 1280)
        self.drop = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(1280, 10)

    def forward(self, x):
        x = self.bn1(nn.functional.relu(self.conv1(x)))
        x = nn.functional.max_pool2d(x, 2)
        x = self.bn2(nn.functional.relu(self.conv2(x)))
        x = nn.functional.max_pool2d(x, 2)
        x = x.view(-1, 64 * 5 * 5)
        x = nn.functional.relu(self.fc1(x))
        x = self.drop(x)
        x = nn.functional.softmax(self.fc2(x), dim=1)
        return x

# Load the data
train_data = torch.load('/content/data/preprocessed/trainset_mnist.pt')
valid_data = torch.load('/content/data/preprocessed/valset_mnist.pt')
test_data = torch.load('/content/data/preprocessed/testset_mnist.pt')

train_dataset = TensorDataset(train_data['data'], train_data['label'])
valid_dataset = TensorDataset(valid_data['data'], valid_data['label'])
test_dataset = TensorDataset(test_data['data'], test_data['label'])

train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
valid_loader = DataLoader(valid_dataset, batch_size=128, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)

# Define the loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.02, momentum=0.9, weight_decay=0.000005)

# Train the model
net = Net()
train_loss_history = []
train_acc_history = []
valid_loss_history = []
valid_acc_history = []

for epoch in range(300):
    train_loss = 0.0
    train_total = 0
    train_correct = 0
    valid_loss = 0.0
    valid_total = 0
    valid_correct = 0
    
    net.train()
    for i, (inputs, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        train_loss += loss.item() * inputs.size(0)
        _, predicted = torch.max(outputs.data, 1)
        train_total += labels.size(0)
        train_correct += (predicted == labels).sum().item()
        
    net.eval()
    with torch.no_grad():
        for inputs, labels in valid_loader:
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            
            valid_loss += loss.item() * inputs.size(0)
            _, predicted = torch.max(outputs.data, 1)
            valid_total += labels.size(0)
            valid_correct += (predicted == labels).sum().item()
    
    train_loss /= len(train_dataset)
    train_acc = train_correct / train_total
    valid_loss /= len(valid_dataset)
    valid_acc = valid_correct / valid_total

train_loss_history.append(train_loss)
train_acc_history.append(train_acc)
valid_loss_history.append(valid_loss)
valid_acc_history.append(valid_acc)

print('Epoch %d: Train Loss: %.4f, Train Acc: %.4f, Valid Loss: %.4f, Valid Acc: %.4f' % (epoch+1, train_loss, train_acc, valid_loss, valid_acc))

#Plot the loss and accuracy history
plt.figure(figsize=(10, 4))
plt.subplot(1, 2, 1)
plt.plot(train_loss_history, label='train')
plt.plot(valid_loss_history, label='valid')
plt.title('Loss History')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(train_acc_history, label='train')
plt.plot(valid_acc_history, label='valid')
plt.title('Accuracy History')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()

# Evaluate the model on the test set
net.eval()
test_total = 0
test_correct = 0
test_pred = []
test_true = []

with torch.no_grad():
  for inputs, labels in test_loader:
    outputs = net(inputs)
    _, predicted = torch.max(outputs.data, 1)
    test_total += labels.size(0)
    test_correct += (predicted == labels).sum().item()
        
    test_pred.extend(predicted.cpu().numpy())
    test_true.extend(labels.cpu().numpy())
    test_acc = test_correct / test_total
    print('Test Acc: %.4f' % test_acc)
    print(classification_report(test_true, test_pred))

# Plot the confusion matrix
cm = confusion_matrix(test_true, test_pred)
plt.figure(figsize=(8, 6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=test_data['label_names'], yticklabels=test_data['label_names'])
plt.title('Confusion Matrix')
plt.xlabel('Predicted Label')
plt.ylabel('True Label')
plt.show()

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix
import seaborn as sns

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class SCNNB(nn.Module):
    def __init__(self):
        super(SCNNB, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(64 * 8 * 8, 1280)
        self.relu3 = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(1280, 10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64 * 8 * 8)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

model = SCNNB().to(device)

batch_size = 64
learning_rate = 0.001
num_epochs = 10

import torch.utils.data as data
import torch

train_data_path = "/content/data/preprocessed/trainset_cifar10.pt"
val_data_path = "/content/data/preprocessed/valset_cifar10.pt"
test_data_path = "/content/data/preprocessed/testset_cifar10.pt"

train_data = torch.load(train_data_path)
val_data = torch.load(val_data_path)
test_data = torch.load(test_data_path)

train_loader = data.DataLoader(train_data, batch_size=batch_size, shuffle=True)
val_loader = data.DataLoader(val_data, batch_size=batch_size, shuffle=False)
test_loader = data.DataLoader(test_data, batch_size=batch_size, shuffle=False)

# model = SCNNB().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

for epoch in range(num_epochs):
    # Training loop
    model.train()
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        if (i+1) % 100 == 0:
            print ('Epoch [{}/{}], Step [{}/{}], Train Loss: {:.4f}' 
                   .format(epoch+1, num_epochs, i+1, len(train_loader), loss.item()))

    # Validation loop
    model.eval()
    with torch.no_grad():
        val_loss = 0
        val_total = 0
        val_correct = 0
        for images, labels in val_loader:
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            val_total += labels.size(0)
            val_correct += (predicted == labels).sum().item()
            val_loss += criterion(outputs, labels).item() * labels.size(0)
        val_loss /= val_total
        val_accuracy = 100 * val_correct / val_total
        print ('Epoch [{}/{}], Validation Loss: {:.4f}, Validation Accuracy: {:.2f}%'.format(epoch+1, num_epochs, val_loss, val_accuracy))

# Test loop
model.eval()
with torch.no_grad():
    test_total = 0
    test_correct = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        test_total += labels.size(0)
        test_correct += (predicted == labels).sum().item()
    test_accuracy = 100 * test_correct / test_total
    print ('Test Accuracy: {:.2f}%'.format(test_accuracy))

import torch.nn as nn

class SCNNB(nn.Module):
    def __init__(self):
        super(SCNNB, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(64 * 5 * 5, 1280)
        self.relu3 = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(1280, 10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64 * 5 * 5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

batch_size = 64
learning_rate = 0.001
num_epochs = 10
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

train_dataset = torch.load('/path/to/train_dataset.pt')
valid_dataset = torch.load('/path/to/valid_dataset.pt')
test_dataset = torch.load('/path/to/test_dataset.pt')

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=batch_size, shuffle=False)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)


# train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        if (i+1) % 100 == 0:
            print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 
                   .format(epoch+1, num_epochs, i+1, len(train_loader), loss.item()))

import torch.nn as nn

class SCNNB(nn.Module):
    def __init__(self):
        super(SCNNB, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(64 * 5 * 5, 1280)
        self.relu3 = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(1280, 10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64 * 5 * 5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

batch_size = 64
learning_rate = 0.001
num_epochs = 10
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        if (i+1) % 100 == 0:
            print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 
                   .format(epoch+1, num_epochs, i+1, len(train_loader), loss.item()))

import torch.nn as nn

class SCNNB(nn.Module):
    def __init__(self):
        super(SCNNB, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(64 * 5 * 5, 1280)
        self.relu3 = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(1280, 10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64 * 5 * 5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

batch_size = 128
learning_rate = 0.001
num_epochs = 2

import torch.utils.data as data
import torch

train_data_path = "/content/data/preprocessed/trainset_cifar10.pt"
val_data_path = "/content/data/preprocessed/valset_cifar10.pt"
test_data_path = "/content/data/preprocessed/testset_cifar10.pt"

train_data = torch.load(train_data_path)
val_data = torch.load(val_data_path)
test_data = torch.load(test_data_path)

train_loader = data.DataLoader(train_data, batch_size=batch_size, shuffle=True)
val_loader = data.DataLoader(val_data, batch_size=batch_size, shuffle=False)
test_loader = data.DataLoader(test_data, batch_size=batch_size, shuffle=False)

model = SCNNB().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

for epoch in range(num_epochs):
    # Training loop
    model.train()
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        if (i+1) % 100 == 0:
            print ('Epoch [{}/{}], Step [{}/{}], Train Loss: {:.4f}' 
                   .format(epoch+1, num_epochs, i+1, len(train_loader), loss.item()))

    # Validation loop
    model.eval()
    with torch.no_grad():
        val_loss = 0
        val_total = 0
        val_correct = 0
        for images, labels in val_loader:
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            val_total += labels.size(0)
            val_correct += (predicted == labels).sum().item()
            val_loss += criterion(outputs, labels).item() * labels.size(0)
        val_loss /= val_total
        val_accuracy = 100 * val_correct / val_total
        print ('Epoch [{}/{}], Validation Loss: {:.4f}, Validation Accuracy: {:.2f}%'.format(epoch+1, num_epochs, val_loss, val_accuracy))

# Test loop
model.eval()
with torch.no_grad():
    test_total = 0
    test_correct = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        test_total += labels.size(0)
        test_correct += (predicted == labels).sum().item()
    test_accuracy = 100 * test_correct / test_total
    print ('Test Accuracy: {:.2f}%'.format(test_accuracy))

import torch
import torch.nn as nn
import torch.optim as optim

class SCNNB(nn.Module):
    def __init__(self):
        super(SCNNB, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(64 * 7 * 7, 1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(1280, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x

model = SCNNB()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.02, momentum=0.9)

import torch
import torch.nn as nn

class SCNN(nn.Module):
    def __init__(self):
        super(SCNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1)
        self.relu2 = nn.ReLU()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.dropout = nn.Dropout(p=0.5)
        self.fc1 = nn.Linear(in_features=64*8*8, out_features=10)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool(x)
        x = x.view(-1, 64*8*8)
        x = self.dropout(x)
        x = self.fc1(x)
        return x
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.transforms as transforms

# Define the hyperparameters
batch_size = 64
learning_rate = 0.001
num_epochs = 20

# Define the transforms to be applied to the data
transform_train = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32, padding=4),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# Load the CIFAR10 dataset
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)

# Create data loaders
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)

# Initialize the model and optimizer
model = SCNN()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# Define the loss function
criterion = nn.CrossEntropyLoss()

# Move the model to the GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Training loop
for epoch in range(num_epochs):
    train_loss = 0.0
    train_correct = 0
    model.train()
    
    for images, labels in train_loader:
        images = images.to(device)
        labels = labels.to(device)
        
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        train_loss += loss.item() * images.size(0)
        _, predicted = torch.max(outputs.data, 1)
        train_correct += (predicted == labels).sum().item()
    
    train_loss /= len(train_loader.dataset)
    train_accuracy = 100.0 * train_correct / len(train_loader.dataset)
    
    # Evaluate the model on the test set
    model.eval()
    test_loss = 0.0
    test_correct = 0
    
    with torch.no_grad():
        for images, labels in test_loader:
            images = images.to(device)
            labels = labels.to(device)
            
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            test_loss += loss.item() * images.size(0)
            _, predicted = torch.max(outputs.data, 1)
            test_correct += (predicted == labels).sum().item()
    
    test_loss /= len(test_loader.dataset)
    test_accuracy = 100.0 * test_correct / len(test_loader.dataset)
    
    # Print the loss and accuracy for each epoch
    print(f'Epoch [{epoch+1}/{num_epochs}], Train Loss: {train_loss:.4f}, Train Accuracy: {train_accuracy:.2f}%, Test Loss: {test_loss:.4f}, Test Accuracy: {test_accuracy:.2f}%')



import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# Load the preprocessed datasets
trainset_mnist = torch.load('./data/preprocessed/trainset_fashion_mnist.pt')
valset_mnist = torch.load('./data/preprocessed/valset_fashion_mnist.pt')
testset_mnist = torch.load('./data/preprocessed/testset_fashion_mnist.pt')

# Define hyperparameters
batch_size = 128
learning_rate = 0.02
momentum = 0.9
weight_decay = 0.000005
num_epochs = 300

# Create data loaders
trainloader_mnist = DataLoader(trainset_mnist, batch_size=batch_size, shuffle=True)
valloader_mnist = DataLoader(valset_mnist, batch_size=batch_size, shuffle=False)
testloader_mnist = DataLoader(testset_mnist, batch_size=batch_size, shuffle=False)

# Define the model
class SCNNB(nn.Module):
    def __init__(self):
        super(SCNNB, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-a
class SCNNB_a(nn.Module):
    def __init__(self):
        super(SCNNB_a, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x

# SCNNB-b
class SCNNB_b(nn.Module):
    def __init__(self):
        super(SCNNB_b, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.relu1 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.relu2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=64*5*5, out_features=1280)
        self.relu3 = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=1280, out_features=10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(-1, 64*5*5)
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x


# Train the models and store the results
models = [SCNNB(), SCNNB_a(), SCNNB_b()]
colors = ['blue', 'red', 'green']
labels = ['SCNNB', 'SCNNB-a', 'SCNNB-b']
train_losses = [[] for _ in range(len(models))]
train_accs = [[] for _ in range(len(models))]
val_losses = [[] for _ in range(len(models))]
val_accs = [[] for _ in range(len(models))]
test_accs = [[] for _ in range(len(models))]

for i, model in enumerate(models):
    optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=weight_decay)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(num_epochs):
        train_loss = 0.0
        train_acc = 0.0
        val_loss = 0.0
        val_acc = 0.0

        # train the model on training set
        model.train()
        for j, data in enumerate(trainloader_mnist, 0):
            inputs, labels = data
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            train_acc += (predicted == labels).sum().item()

        # evaluate the model on validation set
        model.eval()
        with torch.no_grad():
            for k, data in enumerate(valloader_mnist, 0):
                inputs, labels = data
                outputs = model(inputs)
                loss = criterion(outputs, labels)

                val_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                val_acc += (predicted == labels).sum().item()

        # calculate the average loss and accuracy for training and validation sets
        train_loss /= len(trainset_mnist)
        train_acc /= len(trainset_mnist)
        val_loss /= len(valset_mnist)
        val_acc /= len(valset_mnist)

        # save the results
        train_losses[i].append(train_loss)
        train_accs[i].append(train_acc)
        val_losses[i].append(val_loss)
        val_accs[i].append(val_acc)

        # print the loss and accuracy for each epoch
        print(f'Epoch {epoch+1:3d} | Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f} | Test Acc: {test_acc:.4f}')

    # plot the loss and accuracy for training and validation sets
    plt.plot(train_losses[i], label=labels[i] + ' Train Loss', color=colors[i])
    plt.plot(val_losses[i], label=labels[i] + ' Val Loss', linestyle='dashed', color=colors[i])
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.savefig(f"Q21figure1_{labels[i]}.pdf")
    plt.show()

    plt.plot(train_accs[i], label=labels[i] + ' Train Acc', color=colors[i])
    plt.plot(val_accs[i], label=labels[i] + ' Val Acc', linestyle='dashed', color=colors[i])
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.legend()
    plt.savefig(f"Q21figure2_{labels[i]}.pdf")
    plt.show()

    # evaluate the model on test set
    model.eval

# # plot the loss and accuracy for training and validation sets
# import matplotlib.pyplot as plt
# plt.plot(train_losses, label='Train Loss')
# plt.plot(val_losses, label='Val Loss')
# plt.xlabel('Epoch')
# plt.ylabel('Loss')
# plt.legend()
# plt.savefig("Q21figure1.pdf")
# plt.show()

# plt.plot(train_accs, label='Train Acc')
# plt.plot(val_accs, label='Val Acc')
# plt.xlabel('Epoch')
# plt.ylabel('Accuracy')
# plt.legend()
# plt.savefig("Q21figure2.pdf")
# plt.show()

# # plot the accuracy for test set
# plt.plot(test_accs, label='Test Acc')
# plt.xlabel('Epoch')
# plt.ylabel('Accuracy')
# plt.legend()
# plt.savefig("Q21figure3.pdf")
# plt.show()



# evaluate the model on test set
model.eval()
with torch.no_grad():
    test_acc = 0.0
    all_targets = []
    all_predictions = []
    for i, data in enumerate(testloader_mnist, 0):
        inputs, labels = data
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        test_acc += (predicted == labels).sum().item()
        all_targets.extend(labels.cpu().numpy())
        all_predictions.extend(predicted.cpu().numpy())

    # calculate accuracy on test set
    test_acc /= len(testset_mnist)
    test_accs.append(test_acc)
    
    # print classification report
    print('Classification report:')
    print(classification_report(all_targets, all_predictions))
    
    from sklearn.metrics import confusion_matrix
    import seaborn as sns
    import matplotlib.pyplot as plt
    import numpy as np

    # Define the true labels and predicted labels
    true_labels = np.array(test_labels) # assuming `test_labels` contains true labels
    predicted_labels = np.array(test_predictions) # assuming `test_predictions` contains predicted labels

    # Compute confusion matrix
    cm = confusion_matrix(true_labels, predicted_labels)

    # Define class names
    class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

    # Plot confusion matrix
    fig, ax = plt.subplots(figsize=(8, 8))
    sns.heatmap(cm, annot=True, cmap='Blues', cbar=False, xticklabels=class_names, yticklabels=class_names, ax=ax)
    ax.set_xlabel('Predicted labels')
    ax.set_ylabel('True labels')
    ax.set_title('Confusion Matrix')

    # Save and show the plot
    plt.savefig("Q21figure4.pdf")
    plt.show()

"""## Others"""

import torch
import torchvision
import torchvision.transforms as transforms
import numpy as np
import matplotlib.pyplot as plt

# CIFAR-10 dataset
transform_cifar = transforms.Compose([
    transforms.Grayscale(),
    transforms.Resize(28),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

trainset_cifar = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_cifar)
testset_cifar = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_cifar)

# MNIST dataset
transform_mnist = transforms.Compose([
    transforms.Grayscale(),
    transforms.Resize(28),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

trainset_mnist = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform_mnist)
testset_mnist = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform_mnist)

# Fashion-MNIST dataset
transform_fashion = transforms.Compose([
    transforms.Grayscale(),
    transforms.Resize(28),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

trainset_fashion = torchvision.datasets.FashionMNIST(root='./data', train=True, download=True, transform=transform_fashion)
testset_fashion = torchvision.datasets.FashionMNIST(root='./data', train=False, download=True, transform=transform_fashion)

# CIFAR-10 dataset
trainset_cifar, valset_cifar = torch.utils.data.random_split(trainset_cifar, [40000, 10000])
trainloader_cifar = torch.utils.data.DataLoader(trainset_cifar, batch_size=64, shuffle=True)
valloader_cifar = torch.utils.data.DataLoader(valset_cifar, batch_size=64, shuffle=True)
testloader_cifar = torch.utils.data.DataLoader(testset_cifar, batch_size=64, shuffle=False)

# MNIST dataset
trainset_mnist, valset_mnist = torch.utils.data.random_split(trainset_mnist, [50000, 10000])
trainloader_mnist = torch.utils.data.DataLoader(trainset_mnist, batch_size=64, shuffle=True)
valloader_mnist = torch.utils.data.DataLoader(valset_mnist, batch_size=64, shuffle=True)
testloader_mnist = torch.utils.data.DataLoader(testset_mnist, batch_size=64, shuffle=False)

# Fashion-MNIST dataset
trainset_fashion, valset_fashion = torch.utils.data.random_split(trainset_fashion, [50000, 10000])
trainloader_fashion = torch.utils.data.DataLoader(trainset_fashion, batch_size=64, shuffle=True)
valloader_fashion = torch.utils.data.DataLoader(valset_fashion, batch_size=64, shuffle=True)
testloader_fashion = torch.utils.data.DataLoader(testset_fashion, batch_size=64, shuffle=False)

# Show one-picture example for each dataset before and after pre-processing
def show_example(dataset_name, loader):
    # Get one batch of images
    images, labels = next(iter(loader))
    # Denormalize images
    images = images * 0.5 + 0.5
    # Convert images to numpy arrays
    images = images.numpy()

    # Plot the images
    fig, axs = plt.subplots(1, 2, figsize=(5,5))
    axs[0].imshow(np.transpose(images[0], (1, 2, 0)))
    axs[0].set_title('Before pre-processing')
    axs[0].axis('off')
    axs[1].imshow(images[0, 0], cmap='gray')
    axs[1].set_title('After pre-processing')
    axs[1].axis('off')
    fig.suptitle(dataset_name, fontsize=16)
    plt.show()

# Show one example for each dataset before and after pre-processing
show_example('CIFAR-10', trainloader_cifar)
show_example('MNIST', trainloader_mnist)
show_example('Fashion-MNIST', trainloader_fashion)

# Save all the results in a plot.pdf
with PdfPages('plot.pdf') as pdf:
    pdf.savefig(plt.figure(1))
    pdf.savefig(plt.figure(2))
    pdf.savefig(plt.figure(3))

"""## Others"""

import torch
import torchvision
import torchvision.transforms as transforms
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

# Define transforms for pre-processing
transform_gray = transforms.Compose([
    transforms.Grayscale(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5])
])

transform_color = transforms.Compose([
    transforms.Resize((28, 28)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

# Load datasets
cifar10_train = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_color)
mnist_train = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform_gray)
fashion_train = torchvision.datasets.FashionMNIST(root='./data', train=True, download=True, transform=transform_gray)

# Resize test datasets to 28x28
transform_resize = transforms.Compose([
    transforms.Resize((28, 28)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5])
])

cifar10_test = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_resize)
mnist_test = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform_resize)
fashion_test = torchvision.datasets.FashionMNIST(root='./data', train=False, download=True, transform=transform_resize)

# Randomly flip images in Fashion-MNIST dataset
transform_flip = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5])
])

fashion_train_flip = torchvision.datasets.FashionMNIST(root='./data', train=True, download=True, transform=transform_flip)
fashion_test_flip = torchvision.datasets.FashionMNIST(root='./data', train=False, download=True, transform=transform_flip)

# Plot example images
def plot_examples(dataset, title):
    fig, axes = plt.subplots(nrows=2, ncols=5, figsize=(10, 4))
    fig.suptitle(title, fontsize=16)
    for i, ax in enumerate(axes.flat):
        img, _ = dataset[i]
        img = img.numpy().transpose((1, 2, 0))  # Convert to HWC format
        img = (img * 0.5) + 0.5  # Unnormalize
        ax.imshow(img)
        ax.axis('off')
    plt.tight_layout()
    plt.show()
    plt.savefig(title + '.pdf')

# Plot example images before and after pre-processing
plot_examples(cifar10_train, 'CIFAR-10 (color) before pre-processing')
plot_examples(mnist_train, 'MNIST (gray) before pre-processing')
plot_examples(fashion_train, 'Fashion-MNIST (gray) before pre-processing')

plot_examples(cifar10_test, 'CIFAR-10 (color) after pre-processing')
plot_examples(mnist_test, 'MNIST (gray) after pre-processing')
plot_examples(fashion_test, 'Fashion-MNIST (gray) after pre-processing')

plot_examples(fashion_train_flip, 'Fashion-MNIST (gray, flip) before pre-processing')
plot_examples(fashion_test_flip, 'Fashion-MNIST (gray, flip) after pre-processing')

import torch
import torchvision
import torchvision.transforms as transforms
import numpy as np
import matplotlib.pyplot as plt

# Set random seed for reproducibility
torch.manual_seed(42)

# Define the transforms to apply to the datasets
cifar_transforms = transforms.Compose([
    transforms.Grayscale(),
    transforms.Resize(28),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

mnist_transforms = transforms.Compose([
    transforms.Grayscale(),
    transforms.Resize(28),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

fashion_mnist_transforms = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.Grayscale(),
    transforms.Resize(28),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# Load the datasets
cifar_trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=cifar_transforms)

cifar_testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=cifar_transforms)

mnist_trainset = torchvision.datasets.MNIST(root='./data', train=True,
                                        download=True, transform=mnist_transforms)

mnist_testset = torchvision.datasets.MNIST(root='./data', train=False,
                                       download=True, transform=mnist_transforms)

fashion_mnist_trainset = torchvision.datasets.FashionMNIST(root='./data', train=True,
                                        download=True, transform=fashion_mnist_transforms)

fashion_mnist_testset = torchvision.datasets.FashionMNIST(root='./data', train=False,
                                       download=True, transform=fashion_mnist_transforms)

# Split the CIFAR-10 trainset into train and validation sets
cifar_trainset, cifar_valset = torch.utils.data.random_split(cifar_trainset, [40000, 10000])

# Save the pre-processed datasets as train, validation, and test sets separately
torch.save(cifar_trainset, './data/cifar_trainset.pt')
torch.save(cifar_valset, './data/cifar_valset.pt')
torch.save(cifar_testset, './data/cifar_testset.pt')
torch.save(mnist_trainset, './data/mnist_trainset.pt')
torch.save(mnist_testset, './data/mnist_testset.pt')
torch.save(fashion_mnist_trainset, './data/fashion_mnist_trainset.pt')
torch.save(fashion_mnist_testset, './data/fashion_mnist_testset.pt')

# Plot one example image from each dataset before and after pre-processing
cifar_example = cifar_trainset[0][0].squeeze().numpy()
cifar_example = (cifar_example * 0.5) + 0.5  # unnormalize
mnist_example = mnist_trainset[0][0].squeeze().numpy()
mnist_example = (mnist_example * 0.5) + 0.5  # unnormalize
fashion_mnist_example = fashion_mnist_trainset[0][0].squeeze().numpy()
fashion_mnist_example = (fashion_mnist_example * 0.5) + 0.5  # unnormalize
fashion_mnist_flipped_example = fashion_mnist_transforms(fashion_mnist_trainset[0][0])[0].squeeze().numpy()
fashion_mnist_flipped_example = (fashion_mnist_flipped_example * 0.5) + 0.5 # unnormalize

fig, axs = plt.subplots(2, 4, figsize=(12, 6))

axs[0, 0].imshow(np.transpose(cifar_example, (1, 2, 0)), cmap='gray')
axs[0, 0].set_title('CIFAR-10\nBefore preprocessing')
axs[0, 1].imshow(np.transpose(mnist_example, (1, 2, 0)), cmap='gray')
axs[0, 1].set_title('MNIST\nBefore preprocessing')
axs[0, 2].imshow(np.transpose(fashion_mnist_example, (1, 2, 0)), cmap='gray')
axs[0, 2].set_title('Fashion-MNIST\nBefore preprocessing')
axs[0, 3].imshow(np.transpose(fashion_mnist_flipped_example, (1, 2, 0)), cmap='gray')
axs[0, 3].set_title('Fashion-MNIST\nAfter flipping')

cifar_example = cifar_transforms(cifar_trainset[0][0])
cifar_example = (cifar_example.squeeze().numpy() * 0.5) + 0.5 # unnormalize
mnist_example = mnist_transforms(mnist_trainset[0][0])
mnist_example = (mnist_example.squeeze().numpy() * 0.5) + 0.5 # unnormalize
fashion_mnist_example = fashion_mnist_transforms(fashion_mnist_trainset[0][0])
fashion_mnist_example = (fashion_mnist_example.squeeze().numpy() * 0.5) + 0.5 # unnormalize
fashion_mnist_flipped_example = fashion_mnist_transforms(fashion_mnist_trainset[0][0])[0].squeeze().numpy()
fashion_mnist_flipped_example = (fashion_mnist_flipped_example * 0.5) + 0.5 # unnormalize

axs[1, 0].imshow(np.transpose(cifar_example, (1, 2, 0)), cmap='gray')
axs[1, 0].set_title('CIFAR-10\nAfter preprocessing')
axs[1, 1].imshow(np.transpose(mnist_example, (1, 2, 0)), cmap='gray')
axs[1, 1].set_title('MNIST\nAfter preprocessing')
axs[1, 2].imshow(np.transpose(fashion_mnist_example, (1, 2, 0)), cmap='gray')
axs[1, 2].set_title('Fashion-MNIST\nAfter preprocessing')
axs[1, 3].imshow(np.transpose(fashion_mnist_flipped_example, (1, 2, 0)), cmap='gray')
axs[1, 3].set_title('Fashion-MNIST\nAfter flipping and preprocessing')

for ax in axs.flat:
  ax.axis('off')

  plt.savefig('plot.pdf', bbox_inches='tight')
  plt.show()

import torch
import torchvision
import torchvision.transforms as transforms
import random
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt


# Set the seed for reproducibility
torch.manual_seed(0)
random.seed(0)
np.random.seed(0)


# Define the transforms for the datasets
cifar10_transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=1),
    transforms.Resize((28, 28)),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5,), std=(0.5,))
])

mnist_transform = transforms.Compose([
    transforms.Resize((28, 28)),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5,), std=(0.5,))
])

fashion_mnist_transform = transforms.Compose([
    transforms.Resize((28, 28)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5,), std=(0.5,))
])


# Load the datasets with the pre-defined transforms
cifar10_train = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=cifar10_transform)
cifar10_test = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=cifar10_transform)

mnist_train = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=mnist_transform)
mnist_test = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=mnist_transform)

fashion_mnist_train = torchvision.datasets.FashionMNIST(root='./data', train=True, download=True, transform=fashion_mnist_transform)
fashion_mnist_test = torchvision.datasets.FashionMNIST(root='./data', train=False, download=True, transform=fashion_mnist_transform)


# Create the validation sets by splitting the training sets
cifar10_train_size = len(cifar10_train)
cifar10_val_size = int(cifar10_train_size * 0.1)
cifar10_train_size -= cifar10_val_size
cifar10_train, cifar10_val = torch.utils.data.random_split(cifar10_train, [cifar10_train_size, cifar10_val_size])

mnist_train_size = len(mnist_train)
mnist_val_size = int(mnist_train_size * 0.1)
mnist_train_size -= mnist_val_size
mnist_train, mnist_val = torch.utils.data.random_split(mnist_train, [mnist_train_size, mnist_val_size])

fashion_mnist_train_size = len(fashion_mnist_train)
fashion_mnist_val_size = int(fashion_mnist_train_size * 0.1)
fashion_mnist_train_size -= fashion_mnist_val_size
fashion_mnist_train, fashion_mnist_val = torch.utils.data.random_split(fashion_mnist_train, [fashion_mnist_train_size, fashion_mnist_val_size])


# Save the pre-processed datasets as train, validation, and test sets
torch.save(cifar10_train, './data/cifar10_train.pth')
torch.save(cifar10_val, './data/cifar10_val.pth')
torch.save(cifar10_test, './data/cifar10_test.pth')

torch.save(mnist_train, './data/mnist_train.pth')
torch.save(mnist_val, './data/mnist_val.pth')
torch.save(mnist_test, './data/mnist_test.pth')

torch.save(fashion_mnist_train, './data/fashion_mnist_train.pth')
torch.save(fashion_mnist_val, './data/fashion_mnist_val.pth')
torch.save(fashion_mnist_test, './data/fashion_mnist_test.pth')


# Show one example for each dataset before and after the pre-processing steps
def show_example(dataset_name, dataset, index):
    # Get the example image and label from the dataset
    image, label = dataset[index]
    # Convert the image tensor to a numpy array and transpose the dimensions
    image = np.transpose(image.numpy(), (1, 2, 0))
    # Scale the pixel values from [-1, 1] to [0, 1]
    image = (image + 1) / 2
    
    # Check if the image is grayscale and handle it accordingly
    if len(image.shape) == 2:
        image = np.repeat(image[:, :, np.newaxis], 3, axis=2)
    
    # Convert the numpy array to a PIL Image object
    image = Image.fromarray(np.uint8(image * 255))
    
    # Plot the original and pre-processed images side-by-side
    fig, axs = plt.subplots(1, 2, figsize=(5, 2.5))
    fig.suptitle(dataset_name + ' Example #' + str(index))
    
    axs[0].imshow(image)
    axs[0].set_title('Original')
    axs[0].axis('off')
    
    preprocessed_image = dataset[index][0]
    preprocessed_image = np.transpose(preprocessed_image.numpy(), (1, 2, 0))
    preprocessed_image = (preprocessed_image + 1) / 2
    
    # Check if the pre-processed image is grayscale and handle it accordingly
    if len(preprocessed_image.shape) == 2:
        preprocessed_image = np.repeat(preprocessed_image[:, :, np.newaxis], 3, axis=2)
    
    preprocessed_image = Image.fromarray(np.uint8(preprocessed_image * 255))
    
    axs[1].imshow(preprocessed_image)
    axs[1].set_title('Pre-Processed')
    axs[1].axis('off')
    
    plt.savefig('./data/' + dataset_name + '_example_' + str(index) + '.pdf')


show_example('CIFAR-10', cifar10_train, 0)
show_example('MNIST', mnist_train, 0)
show_example('Fashion-MNIST', fashion_mnist_train, 0)

import torch
from torchvision import datasets, transforms

# Define the transforms for data preprocessing
transform = transforms.Compose([
    transforms.Grayscale(), # Convert to grayscale
    transforms.Resize((28, 28)), # Resize to 28x28
    transforms.ToTensor(), # Convert to tensor
    transforms.Normalize((0.5,), (0.5,)) # Normalize to [-1, 1]
])

# Load the CIFAR-10 dataset
cifar10_trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
cifar10_testset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)

# Load the MNIST dataset
mnist_trainset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
mnist_testset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

# Load the Fashion-MNIST dataset
fmnist_trainset = datasets.FashionMNIST(root='./data', train=True, download=True, transform=transform)
fmnist_testset = datasets.FashionMNIST(root='./data', train=False, download=True, transform=transform)

# Flip the training and testing images in Fashion-MNIST randomly
flip_transform = transforms.RandomHorizontalFlip(p=0.5)
fmnist_trainset.transform = transforms.Compose([flip_transform, transform])
fmnist_testset.transform = transform

import torch
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import matplotlib.pyplot as plt

# Define the transforms for each dataset
cifar10_transform = transforms.Compose([
    transforms.Grayscale(),
    transforms.Resize((28, 28, 1)),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])
mnist_transform = transforms.Compose([
    transforms.Grayscale(),
    transforms.Resize((28, 28, 1)),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])
fmnist_transform_train = transforms.Compose([
    transforms.Grayscale(),
    transforms.Resize((28, 28, 1)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])
fmnist_transform_test = transforms.Compose([
    transforms.Grayscale(),
    transforms.Resize((28, 28, 1)),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# Load the datasets
cifar10_trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=cifar10_transform)
cifar10_testset = datasets.CIFAR10(root='./data', train=False, download=True, transform=cifar10_transform)
mnist_trainset = datasets.MNIST(root='./data', train=True, download=True, transform=mnist_transform)
mnist_testset = datasets.MNIST(root='./data', train=False, download=True, transform=mnist_transform)
fmnist_trainset = datasets.FashionMNIST(root='./data', train=True, download=True, transform=fmnist_transform_train)
fmnist_testset = datasets.FashionMNIST(root='./data', train=False, download=True, transform=fmnist_transform_test)

# Get an example from each dataset
cifar10_example, _ = cifar10_trainset[0]
mnist_example, _ = mnist_trainset[0]
fmnist_example, _ = fmnist_trainset[0]

# Create a figure to display the examples
fig, axs = plt.subplots(3, 2, figsize=(6, 9))

# Plot the examples before and after preprocessing
axs[0, 0].imshow(cifar10_example.permute(2, 0, 1))
axs[0, 1].imshow(cifar10_transform(cifar10_example).permute(1, 2, 0))
axs[0, 0].set_title('CIFAR-10 (before)')
axs[0, 1].set_title('CIFAR-10 (after)')

axs[1, 0].imshow(mnist_example.permute(1, 2, 0))
axs[1, 1].imshow(mnist_transform(mnist_example).permute(1, 2, 0))
axs[1, 0].set_title('MNIST (before)')
axs[1, 1].set_title('MNIST (after)')

axs[2, 0].imshow(fmnist_example.squeeze(), cmap='gray')
axs[2, 1].imshow(fmnist_trainset[0][0].squeeze(), cmap='gray')
axs[2, 0].set_title('Fashion-MNIST (before)')
axs[2, 1].set_title('Fashion-MNIST (after)')

# Remove the axes and spacing between subplots
for ax in axs.flat:
    ax.axis('off')
plt.subplots_adjust(wspace=0, hspace=0)

# Save the figure as a PDF
with PdfPages('preprocessing_examples.pdf') as pdf:
    pdf.savefig(fig)

import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages

# Get one example from each dataset
cifar10_example = cifar10_trainset[0][0].squeeze()
mnist_example = mnist_trainset[0][0].squeeze()
fmnist_example = fmnist_trainset[0][0].squeeze()

# Create a figure with two columns and three rows
fig, axs = plt.subplots(nrows=3, ncols=2, figsize=(6, 9))

# Plot the examples before and after preprocessing
axs[0, 0].imshow(cifar10_example.permute(1, 2, 0))
axs[0, 1].imshow(transform(cifar10_example).permute(1, 2, 0))
axs[0, 0].set_title('CIFAR-10 (before)')
axs[0, 1].set_title('CIFAR-10 (after)')

axs[1, 0].imshow(mnist_example, cmap='gray')
axs[1, 1].imshow(transform(mnist_example).squeeze(), cmap='gray')
axs[1, 0].set_title('MNIST (before)')
axs[1, 1].set_title('MNIST (after)')

axs[2, 0].imshow(fmnist_example, cmap='gray')
axs[2, 1].imshow(fmnist_trainset[0][0].squeeze(), cmap='gray')
axs[2, 0].set_title('Fashion-MNIST (before)')
axs[2, 1].set_title('Fashion-MNIST (after)')

# Remove the axes and spacing between subplots
for ax in axs.flat:
    ax.axis('off')
plt.subplots_adjust(wspace=0, hspace=0)

# Save the figure as a PDF
with PdfPages('preprocessing_examples.pdf') as pdf:
    pdf.savefig(fig)

import matplotlib.pyplot as plt

# Get one example from each dataset
cifar10_example = cifar10_trainset[0][0].squeeze()
mnist_example = mnist_trainset[0][0].squeeze()
fmnist_example = fmnist_trainset[0][0].squeeze()

# Create a figure with two columns and three rows
fig, axs = plt.subplots(nrows=3, ncols=2, figsize=(6, 9))

# Plot the examples before and after preprocessing
axs[0, 0].imshow(cifar10_example.permute(1, 2, 0))
axs[0, 1].imshow(transform(cifar10_example).permute(1, 2, 0))
axs[0, 0].set_title('CIFAR-10 (before)')
axs[0, 1].set_title('CIFAR-10 (after)')

axs[1, 0].imshow(mnist_example, cmap='gray')
axs[1, 1].imshow(transform(mnist_example).squeeze(), cmap='gray')
axs[1, 0].set_title('MNIST (before)')
axs[1, 1].set_title('MNIST (after)')

axs[2, 0].imshow(fmnist_example, cmap='gray')
axs[2, 1].imshow(transform(fmnist_example).squeeze(), cmap='gray')
axs[2, 0].set_title('Fashion-MNIST (before)')
axs[2, 1].set_title('Fashion-MNIST (after)')

# Remove ticks from the plots
for ax in axs.flatten():
    ax.set_xticks([])
    ax.set_yticks([])

# Adjust the spacing between the subplots
plt.subplots_adjust(wspace=0.05, hspace=0.3)

# Save the figure as
plt.savefig('preprocessing_examples.pdf')