Training a CNN on MNIST

In this meeting we train a small convolutional neural network to recognize handwritten digits. The walkthrough uses PyTorch and the MNIST-in-CSV dataset on Kaggle.

Imports

Pull in pandas for the CSV, PyTorch for tensors and the model, and matplotlib to peek at digits.

In [1]: Imports

import pandas as pd
import torch
from torch import nn
from torch.utils.data import TensorDataset, DataLoader, Subset
import matplotlib.pyplot as plt

Load CSV into tensors

Kaggle mounts the dataset under /kaggle/input/datasets/oddrationale/mnist-in-csv/. Each row is a label plus 784 pixel values; we reshape to 1×28×28 and scale to [0, 1].

In [2]: make_data

def make_data(filename):
    path = '/kaggle/input/datasets/oddrationale/mnist-in-csv/'
    table = pd.read_csv(path + filename)
    images = torch.tensor(table.iloc[:, 1:].to_numpy(), dtype=torch.float32)
    images = images.reshape(-1, 1, 28, 28) / 255
    labels = torch.tensor(table.iloc[:, 0].to_numpy(), dtype=torch.long)
    return TensorDataset(images, labels)

Datasets and loaders

Build train/test sets, then take a subset of 8,000 training examples so the demo stays snappy.

In [3]: Loaders

train_data = make_data('mnist_train.csv')
test_data = make_data('mnist_test.csv')
train_loader = DataLoader(Subset(train_data, range(8000)), batch_size=128, shuffle=True)
test_loader = DataLoader(test_data, batch_size=256)
print(len(train_loader.dataset), len(test_data))

Peek at a sample

Confirm a digit looks right before training.

In [4]: Sample image

image, label = train_data[0]
plt.imshow(image.squeeze(), cmap='gray')
plt.title(f'The correct digit is {label.item()}')
plt.axis('off');

Model

One conv layer, ReLU, max-pool, then a linear head to 10 digit classes.

In [5]: CNN

model = nn.Sequential(
    nn.Conv2d(1, 8, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(2),
    nn.Flatten(),
    nn.Linear(8 * 14 * 14, 10)
)
print(model)

Loss and optimizer

Cross-entropy for classification; Adam with a small learning rate.

In [6]: Loss & optimizer

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

Train

Run a short training loop over the batched loaders.

In [7]: Training loop

for epoch in range(10):
    model.train()
    total_loss = 0
    for images, labels in train_loader:
        optimizer.zero_grad()
        loss = loss_fn(model(images), labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(epoch + 1, total_loss / len(train_loader))

Originally range(100); using range(10) so the demo finishes faster on Kaggle.

Test accuracy

Evaluate on the full test set without updating weights.

In [8]: Evaluate

model.eval()
correct = 0
with torch.no_grad():
    for images, labels in test_loader:
        guesses = model(images).argmax(dim=1)
        correct += (guesses == labels).sum().item()
print(f'Test accuracy: {correct / len(test_data):.1%}')

One prediction

Pick a test image and compare the CNN’s guess to the label.

In [9]: Single prediction

image, actual = test_data[0]
with torch.no_grad():
    guess = model(image.unsqueeze(0)).argmax(dim=1).item()
plt.imshow(image.squeeze(), cmap='gray')
plt.title(f'CNN guessed {guess} | Actual digit {actual}')
plt.axis('off');