Build Your First Neural Network in PyTorch to Classify Handwritten Digits

Think back to the last time you tried to get started with machine learning. You searched for a tutorial, found one that opened with a diagram of circles joined by lines, scrolled past three paragraphs about partial derivatives, reached a section titled "The Mathematics of Backpropagation", and quietly closed the tab. Nothing was wrong with you. That tutorial simply chose to explain the engine before letting you drive the car. This one does the opposite. Over the next twenty minutes you will build a working neural network, watch it learn to read handwritten digits, and see it answer correctly on images it has never encountered. The explanations arrive next to the code, exactly when you need them and not a minute earlier.

What a Neural Network Actually Is

Strip away the diagrams and a neural network is a very large arithmetic expression with adjustable numbers in it. You feed numbers in at one end, they get multiplied and added along the way, and numbers come out at the other end. The adjustable numbers are called weights, and learning means changing them until the outputs stop being wrong.

That is genuinely the whole idea. What makes it powerful is scale and repetition. The network you build below holds 109,386 adjustable numbers, and it will change every one of them roughly 2,800 times in fifteen seconds.

The word learning does a lot of work in that sentence, so it is worth being precise about what it means here. The network is never told what a 7 looks like. Nobody writes a rule saying a 7 has a horizontal stroke on top. Instead the network guesses, gets told how wrong the guess was, and adjusts slightly in the direction that would have been less wrong. Repeat that a few thousand times and rules that nobody wrote emerge inside the weights.

Why Classifying Digits Is the Standard First Project

Handwritten digit recognition is the traditional starting point, and not out of nostalgia. It has four properties that make it useful for learning, and understanding those properties tells you what to look for when you pick your own first project later.

Guessing randomly across ten classes gets you about 10%. That number is your floor, and it is the reason the first thing you do below is measure the untrained model rather than start training immediately.

How Training Works, Without the Mathematics

Every training system, from this small network to the largest language models, runs the same four step cycle. Learning the cycle here means you already understand the shape of what happens everywhere else.

The model makes a guess. Data goes in, numbers come out. Early on those numbers are meaningless, because the weights started as noise.

A loss function scores the guess. It compares the guess against the correct answer and produces a single number. Large means badly wrong, near zero means nearly right.

Gradients say which direction helps. For every one of the 109,386 weights, the framework works out whether increasing or decreasing it would have made the loss smaller. You never do this calculation yourself.

An optimizer applies the change. It nudges each weight a small step in the helpful direction. How big that step is comes from the learning rate.

That cycle runs once per batch of images. With 60,000 images in batches of 64, one pass over the data runs it 938 times, and three passes run it 2,814 times.

The Five Pieces Every PyTorch Project Contains

Before writing code, it helps to know the cast. Every PyTorch project you ever read will contain these five, usually in this order, and recognising them turns unfamiliar code into something you can skim.

The rest of this tutorial builds those five in order, then runs them.

What You Are Building

With the concepts in place, here is the concrete task. You are building a classifier that looks at a 28 by 28 pixel greyscale image of a handwritten digit and predicts which digit it is, from 0 to 9. The dataset is MNIST, which holds 60,000 training images and 10,000 test images collected from census workers and high school students.

Install PyTorch and Confirm It Works

Install the CPU build. It is a much smaller download than the CUDA build and everything in this tutorial runs on it. Upgrade pip first, because older versions fail partway through this particular install for a reason worth understanding.

In a terminal:

pip install --upgrade pip
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu

In a Jupyter notebook cell, put a percent sign in front of each line:

%pip install --upgrade pip
%pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu

The percent sign is not decoration. A notebook cell is Python, so pasting a shell command straight in gives you SyntaxError: invalid syntax. A single bare line sometimes survives because the kernel guesses your intent, but two lines never do. The %pip form is also the one that installs into the kernel you are actually running, which matters when a machine has several Python versions. Restart the kernel once the install finishes, then import torch.

Confirm the install before writing any model code. From a terminal:

python -c "import torch, torchvision; print(torch.__version__, torchvision.__version__)"

Or in a notebook cell, where it is ordinary Python and needs no prefix:

import torch, torchvision
print(torch.__version__, torchvision.__version__)

2.13.0+cpu 0.28.0+cpu

If you get a version string, you are ready. Those are the exact versions every output below was produced on.

Step 1: Load the Dataset

torchvision ships dataset loaders that handle downloading and parsing for you. The ToTensor() transform does two jobs: it converts the image into a PyTorch tensor and it rescales pixel values from the 0 to 255 range into 0.0 to 1.0.

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

torch.manual_seed(42)

transform = transforms.ToTensor()
train_data = datasets.MNIST(root="./data", train=True,  download=True, transform=transform)
test_data  = datasets.MNIST(root="./data", train=False, download=True, transform=transform)

print("Training images:", len(train_data))
print("Test images    :", len(test_data))
img, label = train_data[0]
print("Image shape    :", img.shape)
print("Pixel range    :", float(img.min()), "to", float(img.max()))
print("First label    :", label)

Training images: 60000
Test images    : 10000
Image shape    : torch.Size([1, 28, 28])
Pixel range    : 0.0 to 1.0
First label    : 5

Step 2: Batch the Data

You do not feed 60,000 images through the network at once. You feed small groups called batches, which keeps memory use low and makes the weight updates less erratic. DataLoader handles the grouping and the shuffling.

train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
test_loader  = DataLoader(test_data,  batch_size=1000)

batch_imgs, batch_labels = next(iter(train_loader))
print("Batch images shape:", batch_imgs.shape)
print("Batch labels shape:", batch_labels.shape)

Batch images shape: torch.Size([64, 1, 28, 28])
Batch labels shape: torch.Size([64])

Batch size now sits at the front of the shape. That leading dimension is present on almost every tensor that moves through a PyTorch model, and forgetting it causes a large share of beginner errors.

Shuffling is on for training and off for testing. During training, shuffling stops the network from learning the order of the file rather than the content of the images. During testing, order does not affect the score, so there is no reason to shuffle.

Step 3: Define the Network

A PyTorch model is a class that inherits from nn.Module. You declare the layers in __init__ and describe how data flows through them in forward.

class DigitNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.stack = nn.Sequential(
            nn.Linear(28*28, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 10),
        )

    def forward(self, x):
        return self.stack(self.flatten(x))

model = DigitNet()
print(model)
print("Trainable parameters:", f"{sum(p.numel() for p in model.parameters()):,}")

DigitNet(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (stack): Sequential(
    (0): Linear(in_features=784, out_features=128, bias=True)
    (1): ReLU()
    (2): Linear(in_features=128, out_features=64, bias=True)
    (3): ReLU()
    (4): Linear(in_features=64, out_features=10, bias=True)
  )
)
Trainable parameters: 109,386

Two things in that output are worth pausing on. The ten numbers the last layer produces are raw scores called logits, which are not probabilities yet and deliberately so, for reasons the next step explains. And the parameter count of 109,386 is the first number in this tutorial that deserves a proper explanation.

Step 4: Choose a Loss Function and an Optimizer

The loss function scores how wrong a prediction is. The optimizer decides how to change the weights in response.

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

Step 5: Measure Accuracy Before You Train

Establish a baseline. This takes ten seconds and it makes the training result meaningful rather than merely impressive.

def evaluate():
    model.eval()
    correct = 0
    with torch.no_grad():
        for X, y in test_loader:
            correct += (model(X).argmax(1) == y).sum().item()
    return 100 * correct / len(test_data)

print(f"{evaluate():.2f}%")

9.55%

With ten digit classes, random guessing gives about 10%. The untrained network sits right there, which confirms it has learned nothing yet.

Two details in that function matter beyond this tutorial. model.eval() switches layers like dropout and batch normalisation into inference behaviour. torch.no_grad() tells PyTorch to skip building the gradient graph, which makes evaluation faster and uses less memory. Neither is optional in production code.

Step 6: Write the Training Loop

This is the centre of the tutorial. Read the five lines inside the inner loop carefully, because you will write them again in every PyTorch project you ever build.

import time

for epoch in range(1, 4):
    model.train()
    t0 = time.time()
    running = 0.0
    for X, y in train_loader:
        pred = model(X)              # 1. forward pass
        loss = loss_fn(pred, y)      # 2. how wrong were we
        optimizer.zero_grad()        # 3. clear last step's gradients
        loss.backward()              # 4. compute new gradients
        optimizer.step()             # 5. update the weights
        running += loss.item()
    print(f"epoch {epoch} | avg loss {running/len(train_loader):.4f} | "
          f"test accuracy {evaluate():.2f}% | {time.time()-t0:.1f}s")

epoch 1 | avg loss 0.3334 | test accuracy 95.37% | 5.2s
epoch 2 | avg loss 0.1371 | test accuracy 96.49% | 5.3s
epoch 3 | avg loss 0.0939 | test accuracy 96.77% | 5.1s

From 9.55% to 95.37% in a single pass over the data, in about five seconds.

The ordering of steps three and four is the part people get wrong. You clear the gradients before computing new ones, because PyTorch accumulates gradients by default rather than replacing them. The next section makes that behaviour visible.

Notice also that loss keeps falling while accuracy gains shrink. Loss measures confidence across every example, accuracy only counts whether the top guess was right. A model can become substantially more confident while converting very few additional wrong answers into right ones.

Step 7: Make a Single Prediction

Accuracy across 10,000 images is a summary. Looking at one prediction shows you what the model actually emits.

model.eval()
sample_img, sample_label = test_data[0]
with torch.no_grad():
    logits = model(sample_img.unsqueeze(0))
    probs = torch.softmax(logits, dim=1)

print("Logits         :", [round(v, 2) for v in logits.squeeze().tolist()])
print("Predicted digit:", logits.argmax(1).item())
print("Confidence     :", f"{probs.max().item()*100:.2f}%")
print("Actual digit   :", sample_label)

Logits         : [-3.06, -4.34, 0.4, 1.15, -7.02, -3.58, -13.15, 8.68, -3.0, -0.23]
Predicted digit: 7
Confidence     : 99.91%
Actual digit   : 7

The eighth value, 8.68, is far above every other score, and it corresponds to the digit 7. Running softmax converts those raw scores into probabilities that sum to 1, giving 99.91% confidence. The model's second choice was the digit 3 at 0.05%, which is a reasonable confusion since 7 and 3 share an upper stroke.

unsqueeze(0) adds the batch dimension. The model expects a batch, so a single image has to be presented as a batch of one.

Step 8: Save and Load the Model

Training results are worthless if they vanish when the process exits.

torch.save(model.state_dict(), "digitnet.pth")

fresh = DigitNet()
fresh.load_state_dict(torch.load("digitnet.pth"))
fresh.eval()

Saved model size: 430.3 KB

Break It On Purpose

Reading about gradient accumulation is forgettable. Watching it happen is not.

grad magnitude without zero_grad(): [1.2995, 2.5989, 3.8984]

This behaviour is not a design flaw. It is what makes gradient accumulation across several small batches possible when a large batch will not fit in memory, a technique used routinely when training large models. PyTorch simply requires you to opt out of it explicitly.

Error Decoder

These are the errors you are most likely to hit, with the exact text PyTorch prints. Each message below was produced by deliberately writing the broken code.

Where to Go From Here

You now have the full loop: load data, batch it, define a model, choose a loss and optimizer, train, evaluate, predict, save. That skeleton is the same whether the model has 109,386 parameters or 70 billion.

Three directions are worth taking next, in roughly this order.

Replace the linear layers with convolutions. A small convolutional network typically pushes MNIST accuracy past 99% because convolutions understand that nearby pixels are related, which a flattened vector cannot express.

Train on your own images rather than a bundled dataset. This is where most people stall, because it requires writing a custom Dataset class. That is the subject of part four in this series.

Understand the machinery under loss.backward(). You have been calling it without knowing what it does. Part three takes that apart.

Conclusion

The gap between not understanding neural networks and having trained one is a single afternoon, and most of that afternoon is installation. What you built is genuinely a neural network: it has weights, it learned them from data by gradient descent, and it generalises to images it has never seen at 96.77% accuracy.

Only a handful of ideas are worth carrying forward. Shapes matter more than maths, and reading them out loud catches most bugs before you run anything. The five step training loop is universal. Gradients accumulate unless you clear them. Evaluation needs model.eval() and torch.no_grad(). Everything else is variation on those ideas.

Keep the script open and change one thing at a time. Widen the hidden layer to 256 and see whether accuracy moves. Drop the learning rate to 1e-4 and watch training slow down. Remove a ReLU and see how much the network loses. Each of those takes fifteen seconds to test, and that fast feedback loop is the real reason to start on CPU rather than waiting for GPU access.

FAQS