Exercise: Classification problem using the MNIST dataset

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import torchvision
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

import matplotlib.pyplot as plt

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

EPOCHS = 5

data_path = "/mnt/e/gx/stu/dataset/MNIST"

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

trainset = datasets.MNIST(
    root = data_path,
    train = True,
    download = False,
    transform = transform
)

testset = datasets.MNIST(
    root = data_path,
    train = False,
    download = False,
    transform = transform
)

trainloader = DataLoader(
    dataset = trainset,
    batch_size = 64,
    shuffle = True,
    num_workers = 2
)

testloader = DataLoader(
    dataset = testset,
    batch_size = 64,
    shuffle = False,
    num_workers = 2
)

img, label = trainset[0]

print(img.shape)
print(label)

plt.figure(figsize=(2, 2))

plt.imshow(img.squeeze(), cmap='gray')
plt.axis('off')
plt.show()

torch.Size([1, 28, 28])
5

class Net(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1 = nn.Conv2d(1, 6, 5, 1, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(400, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        
        x = self.conv1(x)
        x = F.relu(x)
        x = F.max_pool2d(x, (2, 2))
        
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, (2, 2))

        x = torch.flatten(x, 1)

        x = self.fc1(x)
        x = F.relu(x)

        x = self.fc2(x)
        x = F.relu(x)

        x = self.fc3(x)

        return x
        

net = Net().to(DEVICE)

print(net)

Net(
  (conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
  (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
  (fc1): Linear(in_features=400, out_features=120, bias=True)
  (fc2): Linear(in_features=120, out_features=84, bias=True)
  (fc3): Linear(in_features=84, out_features=10, bias=True)
)

criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(net.parameters(), lr=1e-3)

def train(epoch):
    net.train()
    
    running_loss = 0.0
    
    for i, data in enumerate(trainloader):

        inputs, labels = data
        inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)

        optimizer.zero_grad()
        outputs = net(inputs)
        
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

        if i % 100 == 0:
            print(f'Training turn {i}, the loss is: {loss}')

    print("Training Finished")
    
    return running_loss / len(trainloader)

def test():
    net.eval()

    correct = 0

    with torch.no_grad():
        for i, data in enumerate(testloader):

            inputs, labels = data
            inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)

            outputs = net(inputs)

            loss = criterion(outputs, labels)

            pred = outputs.argmax(dim = 1)

            correct += (pred == labels).sum().item()

    accuracy = correct / len(testset)

    print("Testing Finished")

    return accuracy

if __name__ == '__main__':
    
    best_accuracy = 0.0
    
    for epoch in range(1, EPOCHS+1):
    
        print('=' * 60)
        
        train_loss = train(epoch)
        accuracy = test()
    
        print(f'Epoch {epoch}, loss: {train_loss:.8f}, accuracy: {accuracy:.6f}')
    
    print("FINISH")

============================================================
Training turn 0, the loss is: 2.305048942565918
Training turn 100, the loss is: 0.3461371958255768
Training turn 200, the loss is: 0.3229054808616638
Training turn 300, the loss is: 0.1919938027858734
Training turn 400, the loss is: 0.16382364928722382
Training turn 500, the loss is: 0.1612183153629303
Training turn 600, the loss is: 0.08017416298389435
Training turn 700, the loss is: 0.10204880684614182
Training turn 800, the loss is: 0.11169353127479553
Training turn 900, the loss is: 0.08649619668722153
Training Finished
Testing Finished
Epoch 1, loss: 0.28579614, accuracy: 0.971000
============================================================
Training turn 0, the loss is: 0.011228486895561218
Training turn 100, the loss is: 0.1273617446422577
Training turn 200, the loss is: 0.04727615788578987
Training turn 300, the loss is: 0.16774769127368927
Training turn 400, the loss is: 0.02535524033010006
Training turn 500, the loss is: 0.02557501569390297
Training turn 600, the loss is: 0.041032399982213974
Training turn 700, the loss is: 0.027262777090072632
Training turn 800, the loss is: 0.04279733821749687
Training turn 900, the loss is: 0.06929795444011688
Training Finished
Testing Finished
Epoch 2, loss: 0.07020085, accuracy: 0.984900
============================================================
Training turn 0, the loss is: 0.06425740569829941
Training turn 100, the loss is: 0.010477012023329735
Training turn 200, the loss is: 0.17963530123233795
Training turn 300, the loss is: 0.16028636693954468
Training turn 400, the loss is: 0.0651451051235199
Training turn 500, the loss is: 0.014427585527300835
Training turn 600, the loss is: 0.05993712693452835
Training turn 700, the loss is: 0.006976807489991188
Training turn 800, the loss is: 0.010254672728478909
Training turn 900, the loss is: 0.016815701499581337
Training Finished
Testing Finished
Epoch 3, loss: 0.05020663, accuracy: 0.985100
============================================================
Training turn 0, the loss is: 0.06757818162441254
Training turn 100, the loss is: 0.007288593798875809
Training turn 200, the loss is: 0.015437074936926365
Training turn 300, the loss is: 0.07036122679710388
Training turn 400, the loss is: 0.031115880236029625
Training turn 500, the loss is: 0.02757982164621353
Training turn 600, the loss is: 0.12027733027935028
Training turn 700, the loss is: 0.04546143114566803
Training turn 800, the loss is: 0.10460954159498215
Training turn 900, the loss is: 0.03304055333137512
Training Finished
Testing Finished
Epoch 4, loss: 0.04076795, accuracy: 0.987200
============================================================
Training turn 0, the loss is: 0.08599506318569183
Training turn 100, the loss is: 0.00701539684087038
Training turn 200, the loss is: 0.014266052283346653
Training turn 300, the loss is: 0.001244584214873612
Training turn 400, the loss is: 0.03298366069793701
Training turn 500, the loss is: 0.0017019107472151518
Training turn 600, the loss is: 0.01507837325334549
Training turn 700, the loss is: 0.017205098643898964
Training turn 800, the loss is: 0.004085782449692488
Training turn 900, the loss is: 0.004953817930072546
Training Finished
Testing Finished
Epoch 5, loss: 0.03289216, accuracy: 0.988800
FINISH