Building an MNIST Classifier
in PyTorch

01 — Introduction

A census, a pile of forms, and 500 teenagers

Before there was a benchmark, there was a paperwork problem.

In the late 1980s the US Census Bureau was drowning in handwriting. Every form came back filled in by hand, and every one of them had to be read by a person before it became a number in a table. So the Bureau went to the National Institute of Standards and Technology with a request: build us a collection of handwritten characters large enough to judge whether a machine could plausibly do the reading instead.

NIST spent the next few years gathering samples. Thousands of Census field staff scattered across the country filled out standardised forms. Then, wanting a harder test drawn from a different population, NIST collected from 500 high-school students in Maryland.

The two groups did not write alike. People who fill in forms professionally produce neat, consistent, almost typographic digits. Teenagers produce something else. And in the original release one population sat entirely in the training set while the other sat entirely in the test set — which meant your reported accuracy depended less on your algorithm than on which group you happened to train against.

So the dataset was rebuilt. Thirty thousand images from each group went into training, five thousand from each into testing, and the writers were carefully segregated so that nobody's handwriting appeared on both sides of the split. Each digit was size-normalised and centred inside a 28 × 28 box, the smoothing along the way turning hard black-and-white scans into grayscale.

That rebuilt version is MNIST — Modified NIST — and it is what you will train on here. Sixty thousand training images, ten thousand test images, ten classes, and a question a human answers in a fraction of a second without noticing they answered it: which digit is this?

What this post builds

By the end you will have a classifier that answers that question correctly about 97.7% of the time. The model itself is almost anticlimactic — two linear layers and an activation function, a dozen lines of code. That is rather the point. The interesting part of PyTorch is not the architecture; it is the machinery that surrounds it, and the small number of places where getting the order wrong costs you silently.

So the route runs through the workflow first and arrives at the classifier last:

• Getting data in. Why batches exist at all, and the three tools — transforms, Dataset, DataLoader — that carry a sample from disk to model.
• Defining a model. Two ways to write one, why you call the model instead of its forward method, and what breaks when you skip one line of boilerplate.
• Training it. Five lines in a fixed order, and the three ways to rearrange them that produce no error and no learning.
• Knowing whether it worked. Evaluation mode, and why testing on your training data tells you nothing.
• The mechanics underneath. How loss puts a number on being wrong, and how gradients and optimizers turn that number into an adjustment.
• Running it somewhere fast. Device management — the single most common source of frustrating early errors.
• The build. The classifier end to end, then what actually happened across ten epochs, in charts you can scrub through.

02 — Data Ingestion

Why your data arrives in batches

The first thing PyTorch asks you to change is not your model. It's how the data reaches it.

Imagine the delivery company from module one has grown. Instead of ten deliveries you now have 100,000 records, and the obvious move is to read them all in first and worry about the model later:

# Try to load all delivery data into memory
all_distances = []
all_times = []
for i in range(100000):
    distance, time = load_delivery_record(i)
    all_distances.append(distance)
    all_times.append(time)

Every one of those pieces has to live in your computer's RAM at the same time. At 100,000 records you might be fine. Add a few million rows, or bolt on weather data, traffic patterns and driver info, and the machine runs out of memory and crashes in no time.

So you load the data piece by piece instead. The practical unit is a batch — a smaller, manageable chunk of the full dataset. But batching alone isn't enough. Before data is ready for training it also has to be processed, formatted and served in a shape the model understands, and that is exactly the job of PyTorch's three core data utilities: transforms, Dataset and DataLoader.

03 — Data Prep

The three tools that feed the model

Each one owns a distinct part of the journey from disk to model, and they are designed to hand off to one another.

transforms — clean each sample as it loads

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

Transforms are operations that run on each data point as it is loaded. Compose simply means "do the following things, in this order." The two you will meet first are there for one reason: neural networks are surprisingly picky, and they train much better when every input is a small number centred near zero.

• ToTensor() — converts the data to a PyTorch tensor and scales it into the range 0 to 1. Divide [2, 4, 6, 8, 10] by its max and you get tensor([.2, .4, .6, .8, 1]).
• Normalize() — shifts those values so they centre on zero, then scales them by the standard deviation: tensor([-1.41, -0.70, 0, .70, 1.41]).

Dataset — know where a sample lives and how to fetch it

# PyTorch has many pre-built datasets
dataset = SomeDataset('./data', train=True, download=True, transform=transform)

first_item = dataset[0]          # Just gets one

A dataset fetches each sample from disk when it is asked to, rather than preloading everything in one shot — that is the secret to handling datasets far bigger than memory. It knows four things: where your data lives on disk, how to load a specific sample, how many samples there are in total, and how to apply your transforms to each one on the way out.

The arguments read exactly as they sound. './data' is where the files are stored locally; train=True picks the training split rather than the test split; download=True fetches the data if it isn't already there; and you pull out a single sample by indexing.

DataLoader — serve batches on demand

dataset_loader = DataLoader(dataset, batch_size=32, shuffle=True)

The loader is the part that makes training on large datasets possible: it requests one batch at a time from the dataset. batch_size sets how many samples arrive at once, and shuffle mixes the order, which helps the model learn more effectively during training.

All three together

# 1. Define your transforms
transform = SomeTransform()      # ToTensor + Normalize

# 2. Create Datasets with transforms
train_dataset = YourDataset('./data', train=True, transform=transform)

# 3. Create DataLoaders
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 4. Training
for batch_idx, (data, labels) in enumerate(train_loader):
    # Data arrives in batches, already transformed
    output = model(data)

04 — Modeling

Two ways to define a model

nn.Sequential works beautifully. The other pattern does exactly the same thing with more control — and you will see it everywhere in PyTorch code.

model = nn.Sequential(            class ExampleModel(nn.Module):
    nn.Linear(1, 20),                 def __init__(self):
    nn.ReLU(),                            super().__init__()
    nn.Linear(20, 1)                      self.layer_1 = nn.Linear(1, 20)
)                                         self.relu    = nn.ReLU()
                                          self.layer_2 = nn.Linear(20, 1)

                                      def forward(self, x):
                                          x = self.layer_1(x)
                                          x = self.relu(x)
                                          x = self.layer_2(x)
                                          return x

Every PyTorch module class splits into two halves. __init__ defines which layers exist — a little like gathering your tools. forward describes the path the data actually takes through them. Sequential does the same ordering, just written in a different style.

Call the model, never .forward()

model = ExampleModel()

output = model(data)             # ✔
output = model.forward(data)     # ✘

You wrote a forward method, so calling it directly feels natural. But note that Sequential is itself a subclass of Module with its own forward — and you never called that one by hand either. When you call the model and pass in the data, PyTorch does more than run your method: it makes internal checks, tracks the necessary math, and sets things up for updating the model later. Calling forward yourself skips all of that essential bookkeeping.

And super().__init__() is not boilerplate

class ExampleModel(nn.Module):
    def __init__(self):
        # super().__init__()   # What if you skip this?
        self.fc1 = nn.Linear(784, 128)

# AttributeError: cannot assign parameter before Module.__init__() call
model = ExampleModel()

PyTorch needs somewhere to track all the learnable parameters — the weights and biases you will be updating during training. super().__init__() is what creates that tracking system. Without it, PyTorch has nowhere to register your layers, and it says so immediately.

05 — Training

The training loop, and why the order matters

Five lines. Written in the wrong order they still run, produce no error, and quietly fail to teach the model anything.

for epoch in range(epochs):
    optimizer.zero_grad()                    # 1. Reset the optimizer
    outputs = model(X)                       # 2. Make predictions
    loss = loss_function(outputs, y)         # 3. How bad was this guess?
    loss.backward()                          # 4. Calculate adjustments
    optimizer.step()                         # 5. Update the model

zero_grad clears out the old calculations, backward figures out the improvements, and step applies them. This standard sequence is the core of most PyTorch training loops, and the failure mode when you rearrange it is silence — PyTorch will not throw an error, but your model will not learn properly.

06 — Evaluation

Checking whether it actually learned

Evaluation means testing on data the model has not previously seen and was not trained on. Testing on the training set is like giving a student the same exam twice — they might just have memorised the answers.

model.eval()          # Set evaluation mode (NOT "evaluate my model"!)

with torch.no_grad():  # Disable gradient tracking
    correct = 0
    total = 0
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)   # class with highest score
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f'Accuracy: {accuracy}%')

model.train()         # Set the model back to training mode

• model.eval() — despite the name, this does not evaluate anything. It puts the model into evaluation mode, which is both cheaper computationally and necessary because some layers behave differently during training and evaluation.
• torch.no_grad() — switches off the extra tracking PyTorch does during training. Leave it on and PyTorch keeps storing details it does not need, wasting memory and sometimes crashing the program mid-validation.

For a classification task the score itself is straightforward: count how often the model gets the class right and divide by the total attempts. 9,500 correct out of 10,000 is 95% accuracy. And if you are heading back into training afterwards, remember to switch the mode back.

07 — Measure

Loss: putting a number on how wrong you are

Three lines carry the whole of training, and they run in a fixed sequence — measure, diagnose, update.

loss = loss_function(outputs, targets)   → measure
loss.backward()                          → diagnose
optimizer.step()                         → update

First you measure how wrong the predictions are — one number that sums up all the mistakes. Then backward diagnoses the problem, examining how each weight contributed to that error and by how much. Finally step adjusts each parameter, with the biggest corrections going to the weights that caused the biggest problems.

Every loss function does the same basic job: compare predictions against the true answers and return a number. The higher the number, the more wrong you are.

Mean squared error, for predicting numbers

Back on the delivery data, suppose the model predicts 6 minutes when the truth was 4, and 3 minutes when the truth was 5. Subtract target from prediction and you have the error. Average the raw errors, though, and something absurd happens:

Zero. Perfect performance, apparently, from two predictions that were both wrong. Squaring the differences removes the minus signs so mistakes stop cancelling each other out — and it has a second effect worth having. Being off by 10 minutes scores 100 while being off by 1 scores 1, so a ten-minute miss is a hundred times worse than a one-minute miss.

loss_function = nn.MSELoss()
loss = loss_function(predictions, targets)

So MSE helps in two ways: it makes sure every mistake counts, and it punishes bigger errors more than small ones. It is the right choice whenever you are predicting a continuous value — distances, temperatures, prices.

Cross-entropy, for predicting categories

loss_function = nn.CrossEntropyLoss()
loss = loss_function(outputs, labels)

Here the model isn't picking a single answer — it outputs a confidence score across every possible class. For MNIST that is ten numbers per image, one probability per digit, all adding up to 100%:

Digit:  0     1     2     3     4     5     6     7     8     9
Conf:   0.05  0.02  0.08  0.70  0.03  0.04  0.02  0.03  0.02  0.01

And here is the key idea: cross-entropy punishes overconfident wrong answers. If the model is 95% sure an image is a 7 when it is really a 3, the loss is very high. If it is only 55% sure, it is still wrong but less brazen about it, and the loss is smaller. You want a model that is confident about right answers and unsure about wrong ones — cross-entropy is what shapes that behaviour.

Choosing, and what else is out there

• MSE — predicting a number, such as a temperature, price or distance.
• CrossEntropyLoss — predicting a category, like a digit, animal or word.

nn.MSELoss()            # average squared difference — regression
nn.CrossEntropyLoss()   # multi-class classification (1 out of many)
nn.L1Loss()             # average absolute difference — less outlier-sensitive
nn.BCEWithLogitsLoss()  # binary classification (yes/no, true/false)
nn.NLLLoss()            # multi-class when you're using LogSoftmax
nn.SmoothL1Loss()       # balances MSE and L1 (also called Huber loss)
nn.KLDivLoss()          # difference between two probability distributions

There are plenty more, each suited to a specific situation, but MSE and cross-entropy cover a great deal of territory in deep learning.

08 — Diagnose & Update

Gradients and optimizers

Loss tells you that something is wrong. Gradients tell you which weights are responsible.

Each neuron takes its inputs, multiplies each by a weight, adds them up with a bias, and passes the result through an activation function — z = W1*x1 + W2*x2 + W3*x3 + b. Even a small network involves a startling number of those:

backward works through all of them like a detective, asking each weight and bias the same question: how much did you contribute to the loss? The answers are the gradients.

Downhill, one step at a time

Minimising loss is like standing on a hillside trying to reach the bottom of a valley. The gradient tells you the slope where you're standing — which way is up, which way is down. You head downhill, toward lower loss. Hence gradient descent.

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

Stochastic gradient descent has a simple strategy: negative gradient, increase the weight; positive gradient, decrease it; big gradient, big change; small gradient, small change. It does not subtract the gradient directly, though — it scales it by the learning rate first, so a gradient of 0.5 with lr=0.01 moves the weight by 0.005. That scaling is what makes or breaks training:

• Tiny learning rate — tiny steps; it will take forever to reach the bottom.
• Good learning rate — steady progress all the way down.
• Huge learning rate — giant leaps that bounce back and forth, overshooting the minimum.

# Smaller learning rate
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
optimizer = torch.optim.SGD(model.parameters(),  lr=0.01)

SGD works well, but smarter optimizers adapt to each weight individually. Adam is one of those — something like an assistant who knows which weights need big adjustments and which need fine-tuning. It has become a popular first choice because it is reliable, flexible and often faster than the alternatives. One word of caution: do not copy the learning rate from SGD when you switch to Adam, because it is tuned completely differently and your loss might explode. Beyond these two sit RMSprop, Adagrad, AdamW, NAdam and a dozen more, but for most projects SGD and Adam have you covered.

Which finally explains zero_grad

Every time you call backward, PyTorch adds the new gradients to whatever is already sitting there. Skip zero_grad and you are no longer diagnosing this batch — you are accumulating the diagnoses of every batch so far, and they keep piling up incorrectly until training breaks. That accumulation is deliberate, because it enables advanced tricks like gradient accumulation and custom training schedules, but for ordinary work you want a clean slate at the top of every loop.

09 — Hardware

Device management

Every tensor and every model lives on a device. PyTorch will not move things around for you automatically, and if they are not all in the same place your code may simply refuse to run.

RuntimeError: Expected all tensors to be on the same device

A CPU is built for general-purpose work and runs operations sequentially. A GPU runs them in parallel, and during training that can be 10 to 15 times faster. If your system has one, you almost always want to use it — you just have to place things by hand.

# 1. Is there an accelerator?
torch.cuda.is_available()        # True → GPU available

# 2. Pick a device — the safe default you'll see everywhere
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 3. Move the model once, when you create it
model = MyModel().to(device)

# 4. Move each batch, every iteration
for inputs, targets in dataloader:
    inputs = inputs.to(device)
    targets = targets.to(device)

cuda refers to NVIDIA GPUs and their toolkit of the same name. There are other options — MPS for Apple Silicon, for instance — but CUDA is the most commonly used.

Finding out where something lives

# For tensors
print(inputs.device)

# For models
print(next(model.parameters()).device)

Models themselves aren't on a device — their parameters are, so you check one parameter to learn where they all are. If you are still seeing device errors after that, check your targets and your model's outputs too.

x.to(device)        # looks like it moves x — but it doesn't
x = x.to(device)    # this does

The whole loop, placed correctly

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyModel().to(device)
optimizer = optim.Adam(model.parameters())
loss_function = nn.CrossEntropyLoss()

for inputs, targets in dataloader:
    inputs = inputs.to(device)
    targets = targets.to(device)
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss.backward()
    optimizer.step()

Three steps, and they are the foundation of every training script you will write: choose your device up front, move the model once, move the data every batch.

And when the GPU runs out of room

GPU memory is limited. If your model and batch size need more than it has, you get CUDA out of memory. Small batches make training slow; batches that are too large crash it. For many systems a batch size between 32 and 64 is a good starting point, though it depends on your hardware and architecture. If you see a memory error, lower the batch size first — it is the most common fix.

10 — Putting it together

An MNIST classifier, start to finish

60,000 training images, 10,000 test images, each one 28 × 28 pixels of grayscale handwriting, centred, digits 0 through 9.

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

TorchVision is PyTorch's computer vision library. It ships with the popular datasets, MNIST among them, plus the image-processing tools.

Step 1 — preprocessing

# Data preprocessing
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

ToTensor scales the pixels from 0–255 down to 0–1, and Normalize shifts and scales those values so they centre on zero. Those two odd-looking numbers are the mean and standard deviation of the entire MNIST training set — normalising every image by the same pair keeps the data consistent, which helps the model learn faster.

Step 2 — the datasets

# Load MNIST dataset
train_dataset = torchvision.datasets.MNIST(
    root='./data', train=True, download=True, transform=transform)

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

The test set is almost identical — the single change is train=False, which gives you the 10,000 test images instead of the 60,000 training ones. Same transforms, same storage location; TorchVision handles the downloading and organising.

Step 3 — the loaders, and why only one of them shuffles

# Create data loaders
train_loader = DataLoader(train_dataset, batch_size=64,   shuffle=True)
test_loader  = DataLoader(test_dataset,  batch_size=1000, shuffle=False)

Training runs 64 images per batch, reshuffled every epoch so the model meets them in a different random order each time. Testing uses much larger batches — there are no gradients to compute, so there is no reason to go in small.

The asymmetry is worth pausing on. Datasets often arrive organised by class, so without shuffling your model might see 6,000 zeros in a row before it meets a single one. It could learn an unintended pattern — early batches are zeros, late batches are nines — instead of learning what actually makes a zero look like a zero. Shuffling mixes everything so each batch carries variety. At test time the model has finished learning and you are only checking whether it recognises digits, so the order stops mattering.

Step 4 — the network

class MNISTClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.layers = nn.Sequential(
            nn.Linear(784, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )

    def forward(self, x):
        x = self.flatten(x)
        x = self.layers(x)
        return x

Flatten is the new piece, and it exists because of a shape mismatch. A single MNIST image arrives as a [1, 28, 28] tensor — one channel (grayscale, a single brightness value per pixel) by 28 by 28. Batch it up and PyTorch adds a dimension, so at batch_size=64 the data reaches your model as [64, 1, 28, 28].

Linear layers expect flat vectors: one long row of numbers per image, not a two-dimensional grid. Flatten reshapes each 28 × 28 image into 784 values in a row — 28 × 28 = 784 — turning [64, 1, 28, 28] into [64, 784]. Leave it out and the image data hits the linear layer and raises a shape mismatch error.

After that the stack is short: 784 pixel values become 128 hidden features, ReLU keeps the positives and zeroes the negatives, and 128 features become 10 outputs — one per digit class.

Step 5 — device, loss, optimizer

# Check for GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using {device}')

# Initialize model and move to device
model = MNISTClassifier().to(device)

# Loss function and optimizer
loss_function = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

Cross-entropy because the model picks one class out of many, which is exactly the job of choosing a digit from 0 to 9. Adam because it adapts its learning rate as it goes — larger adjustments early, when gradients are noisy, and smaller corrections later as training stabilises.

Step 6 — one epoch of training

def train_epoch(model, train_loader, loss_function, optimizer, device):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = loss_function(output, target)
        loss.backward()
        optimizer.step()

        # Track progress
        running_loss += loss.item()
        _, predicted = output.max(1)
        total += target.size(0)
        correct += predicted.eq(target).sum().item()

Five inputs — the model, the loader, the loss function, the optimizer, and the device everything should run on. model.train() puts the model into training mode, and three counters track the epoch: running_loss accumulates loss values, correct counts predictions matching the true labels, total counts the samples seen. Inside the loop it is the familiar sequence, bookended by moving the batch onto the device and recording what happened. output.max(1) is what tells you which digit class scored highest.

With 60,000 images at batch size 64, an epoch is about 938 batches — so printing every hundredth gives you roughly nine progress lines per pass.

Step 7 — evaluation, then ten epochs

def evaluate(model, test_loader, device):
    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for inputs, targets in test_loader:
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()

    return 100. * correct / total


# Training loop
num_epochs = 10
for epoch in range(num_epochs):
    print(f'\nEpoch: {epoch+1}')
    train_epoch(model, train_loader, loss_function, optimizer, device)
    accuracy = evaluate(model, test_loader, device)
    print(f'Test Accuracy: {accuracy:.2f}%')

The evaluation function is the training function with everything unnecessary removed: no optimizer, no loss tracking, no weight updates. Just how many did you get right.

Ten epochs means ten full passes over the training set, and it is not mere repetition — each pass refines the model's sense of what makes a 2 different from a 7. Evaluating on the test set after every epoch is what tells you whether it is learning generalisable patterns or simply memorising.

11 — Results

What the training actually looked like

Two views of the same ten epochs: the curves, and the digits themselves.

The loss and accuracy history is the standard instrument panel. The left panel tracks loss per epoch for both splits, the right tracks accuracy, and the wide panel underneath plots the loss of every single one of the 9,380 batches — the noise the epoch averages smooth away.

Training history. Train loss falls from 0.254 to 0.020 and train accuracy climbs from 92.5% to 99.3%. Test accuracy, though, flattens out around 97.5% after the fourth epoch — and test loss bottoms out at 0.081 on epoch 8 before drifting back up to 0.094 while train loss keeps falling. That gap opening up is the model beginning to memorise the training set. It is the clearest argument for evaluating on held-out data every epoch rather than trusting the training numbers.

The curves tell you the model improved. They do not show you what improving looks like. For that, ten test digits — one per class — were pulled out before training started and re-predicted after every epoch, so you can watch the labels change. Press play, or drag the slider.

Actual vs predicted, epoch by epoch. Green means the prediction matches the true label, red means it doesn't, and the percentage is how confident the model was. Before training the model scores 15.99% and its confidences hover around 12–19% — barely distinguishable from picking at random across ten classes. One epoch later it has nine of the ten right and most confidences are above 95%. The stubborn cases are the interesting ones: the 5 is read as a 6 at 68% confidence after epoch one and is not corrected until epoch four, and even at epoch ten the 3 is only held at 70%. Ambiguous handwriting stays ambiguous.