text_plain_output_1.png
from torch.utils.data import Dataset, DataLoader

from torchvision import transforms

import numpy as np

import random

import torch

import torch.nn as nn

import torch.optim as optim

import torch.optim.lr_scheduler as lr_scheduler

import torchvision

def set_seed(seed=0):
    np.random.seed(seed)
    random.seed(seed)
    torch.manual_seed(seed)
set_seed()

BATCH_SIZE = 64
LEARNING_RATE = 0.01
N_EPOCHS = 20

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

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

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

train_loader = DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=BATCH_SIZE, shuffle=False)

classes = ['plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']

class ConvNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=5, stride=1, padding=0)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1, padding=0)
        self.flat = nn.Flatten()
        self.fc1 = nn.Linear(in_features=32 * 5 * 5, out_features=128)
        self.fc2 = nn.Linear(in_features=128, out_features=64)
        self.fc3 = nn.Linear(in_features=64, out_features=10)
        self.relu = nn.ReLU()
    def forward(self, x):
        out = self.conv1(x)
        out = self.relu(out)
        out = self.pool(out)
        out = self.conv2(out)
        out = self.relu(out)
        out = self.pool(out)
        out = self.flat(out)
        out = self.fc1(out)
        out = self.relu(out)
        out = self.fc2(out)
        out = self.relu(out)
        out = self.fc3(out)
        return out
model = ConvNet().to(device)

loss = nn.CrossEntropyLoss()
optimiser = optim.SGD(params=model.parameters(), lr=LEARNING_RATE)
scheduler = lr_scheduler.CosineAnnealingLR(optimiser, T_max=N_EPOCHS)

for epoch in range(N_EPOCHS):
    loss_acc = 0
    train_count = 0
    for i, (imgs, labels) in enumerate(train_loader):
        imgs = imgs.to(device)
        labels = labels.to(device)
        preds = model(imgs)
        L = loss(preds, labels)
        L.backward()
        optimiser.step()
        optimiser.zero_grad()
        loss_acc += L.detach().item()
        train_count += 1
    scheduler.step()

PATH = 'ConvNet.pt'
torch.save({'epoch': N_EPOCHS, 'model_state_dict': model.state_dict(), 'optimiser_state_dict': optimiser.state_dict(), 'scheduler_state_dict': scheduler.state_dict(), 'loss': loss_acc / train_count}, PATH)

model = ConvNet().to(device)
optimiser = optim.SGD(params=model.parameters(), lr=LEARNING_RATE)
scheduler = lr_scheduler.CosineAnnealingLR(optimiser, T_max=N_EPOCHS)
checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimiser.load_state_dict(checkpoint['optimiser_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']
model.eval()

with torch.no_grad():
    n_correct = 0
    n_samples = 0
    n_class_correct = [0 for i in range(10)]
    n_class_sample = [0 for i in range(10)]
    for imgs, labels in test_loader:
        imgs = imgs.to(device)
        labels = labels.to(device)
        output = model(imgs)
        _, preds = torch.max(output, 1)
        n_samples += labels.shape[0]
        n_correct += (preds == labels).sum().item()
        for i in range(BATCH_SIZE):
            try:
                label = labels[i].item()
                pred = preds[i].item()
            except:
                break
            if label == pred:
                n_class_correct[label] += 1
            n_class_sample[label] += 1
    acc = 100 * n_correct / n_samples
    print(f'Overall accuracy on test set: {acc:.1f} %')
    for i in range(10):
        print(f'Accuracy of {classes[i]}: {100 * n_class_correct[i] / n_class_sample[i]:.1f} %')