PyTorch MNIST Classifier Example

This PyTorch MNIST classifier example puts everything together: data loading, a model, the training loop, GPU support and evaluation. MNIST contains 70,000 grayscale images of handwritten digits (0-9), each 28 by 28 pixels, split into 60,000 for training and 10,000 for testing.

Step 1: Load the Data

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

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

tf = transforms.Compose([transforms.ToTensor(),
                         transforms.Normalize((0.1307,), (0.3081,))])
train_ds = datasets.MNIST("data", train=True, download=True, transform=tf)
test_ds = datasets.MNIST("data", train=False, download=True, transform=tf)

train_dl = DataLoader(train_ds, batch_size=128, shuffle=True)
test_dl = DataLoader(test_ds, batch_size=512)
print(len(train_ds), len(test_ds))

60000 10000

0.1307 and 0.3081 are the mean and standard deviation of MNIST pixels. Normalizing inputs to roughly zero mean and unit variance helps training converge faster.

Step 2: Define the Model

model = nn.Sequential(
    nn.Flatten(),             # [B, 1, 28, 28] -> [B, 784]
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Dropout(0.2),
    nn.Linear(256, 128),
    nn.ReLU(),
    nn.Linear(128, 10),       # one logit per digit
).to(device)

print(sum(p.numel() for p in model.parameters()))

235146

Step 3: Train and Evaluate

loss_fn = nn.CrossEntropyLoss()
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)

def evaluate():
    model.eval()
    correct = 0
    with torch.no_grad():
        for xb, yb in test_dl:
            xb, yb = xb.to(device), yb.to(device)
            correct += (model(xb).argmax(1) == yb).sum().item()
    return correct / len(test_ds)

for epoch in range(5):
    model.train()
    for xb, yb in train_dl:
        xb, yb = xb.to(device), yb.to(device)
        loss = loss_fn(model(xb), yb)
        opt.zero_grad()
        loss.backward()
        opt.step()
    print(f"epoch {epoch + 1}: last loss {loss.item():.3f}, test acc {evaluate():.4f}")

epoch 1: last loss 0.142, test acc 0.9612
epoch 2: last loss 0.087, test acc 0.9714
epoch 3: last loss 0.061, test acc 0.9755
epoch 4: last loss 0.102, test acc 0.9779
epoch 5: last loss 0.030, test acc 0.9801
(sample output — exact numbers vary per run)

A simple fully connected network reaches around 98% accuracy in a few minutes on a CPU. The single-batch loss jumps around; the test accuracy is the number to watch.

Step 4: Predict a Single Image

model.eval()
img, label = test_ds[0]
with torch.no_grad():
    probs = model(img.unsqueeze(0).to(device)).softmax(dim=1)
print("true:", label, "predicted:", probs.argmax().item(),
      "confidence:", round(probs.max().item(), 3))

true: 7 predicted: 7 confidence: 0.999   (sample output)

unsqueeze(0) adds the batch dimension, because models always expect a batch, even of one.

Ideas to Improve It

• Swap the MLP for a convolutional network (next lesson) to reach about 99%.
• Add a learning rate scheduler.
• Try data augmentation such as small random rotations.
• Try FashionMNIST, a harder drop-in replacement with clothing images.

Frequently asked questions

What accuracy should I expect on MNIST with PyTorch?

A simple fully connected network typically reaches about 97-98%, and a small CNN about 99% on the test set.

Why does MNIST use Normalize((0.1307,), (0.3081,))?

Those are the mean and standard deviation of the training pixels. Normalizing inputs helps gradients behave well and speeds up training.

How do I predict a single image in PyTorch?

Switch to eval mode, add a batch dimension with unsqueeze(0), run the model inside torch.no_grad(), and take argmax of the output.