MNIST Digit Recognition — MLP vs CNN

This project trained two classifiers for the mnist digit dataset. The first uses a standard MLP and the second uses convolutional layers.

Live demo

Draw a digit in the black canvas. Both models predict simultaneously.

Code — Project 3 (MLP)

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.input      = nn.Linear(784, 784)
        self.input_act  = nn.Tanh()
        self.hidden     = nn.Linear(784, 784)
        self.hidden_act = nn.Tanh()
        self.output     = nn.Linear(784, 10)

    def forward(self, x):
        l1 = self.input_act(self.input(x))
        l2 = self.hidden_act(self.hidden(l1))
        return self.output(l2)


if __name__ == "__main__":
    import torch.optim as optim
    import torchvision
    from torch.utils.data import DataLoader
    import os

    transform = torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize((0.5,), (0.5,)),
        torchvision.transforms.Lambda(lambda x: torch.flatten(x)),
    ])

    train_data   = torchvision.datasets.MNIST(root='./data', train=True,  download=True, transform=transform)
    val_data     = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
    train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
    val_loader   = DataLoader(val_data,   batch_size=64)

    model     = Model()
    loss_fn   = nn.CrossEntropyLoss()
    optimizer = optim.AdamW(model.parameters(), lr=0.001)

    for epoch in range(10):
        model.train()
        for inputs, labels in train_loader:
            optimizer.zero_grad()
            loss = loss_fn(model(inputs), labels)
            loss.backward()
            optimizer.step()

        model.eval()
        total_acc, n = 0.0, 0
        with torch.no_grad():
            for inputs, labels in val_loader:
                total_acc += (model(inputs).argmax(dim=1) == labels).float().sum().item()
                n += len(labels)
        print(f"epoch={epoch+1:2d}  val_acc={total_acc/n:.4f}")

    torch.save(model.state_dict(), os.path.join(os.path.dirname(__file__), "model.pt"))

Code — Project 4 (CNN)

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1      = nn.Conv2d(1, 32, 3)   # 28x28 → 26x26
        self.conv1_act  = nn.ReLU()
        self.conv1_pool = nn.MaxPool2d(2)        # → 13x13
        self.conv2      = nn.Conv2d(32, 64, 3)  # → 11x11
        self.conv2_act  = nn.ReLU()
        self.conv2_pool = nn.MaxPool2d(2)        # → 5x5
        self.output     = nn.Linear(1600, 10)   # 64×5×5 = 1600

    def forward(self, x):
        x = self.conv1_pool(self.conv1_act(self.conv1(x)))
        x = self.conv2_pool(self.conv2_act(self.conv2(x)))
        return self.output(torch.flatten(x, 1))


if __name__ == "__main__":
    import torch.optim as optim
    import torchvision
    from torch.utils.data import DataLoader
    import os

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

    train_data   = torchvision.datasets.MNIST(root='./data', train=True,  download=True, transform=transform)
    val_data     = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
    train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
    val_loader   = DataLoader(val_data,   batch_size=64)

    model     = Model()
    loss_fn   = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    for epoch in range(5):
        model.train()
        for inputs, labels in train_loader:
            optimizer.zero_grad()
            loss = loss_fn(model(inputs), labels)
            loss.backward()
            optimizer.step()

        model.eval()
        total_acc, n = 0.0, 0
        with torch.no_grad():
            for inputs, labels in val_loader:
                total_acc += (model(inputs).argmax(dim=1) == labels).float().sum().item()
                n += len(labels)
        print(f"epoch={epoch+1:2d}  val_acc={total_acc/n:.4f}")

    torch.save(model.state_dict(), os.path.join(os.path.dirname(__file__), "model.pt"))