Train an Image Classifier in PyTorch

Build a neural network that recognizes handwritten digits from scratch.

Introduction

30 Second Summary

Every time your phone unlocks with face recognition or your bank reads a handwritten check, a neural network is classifying images behind the scenes. Training one of these networks from scratch is the foundational skill behind all of modern AI.

In this project, you will build and train a neural network in PyTorch that recognizes handwritten digits with over 97% accuracy. You'll write every piece yourself: the model architecture, the data pipeline, the training loop, and the evaluation logic.

What You'll Build

A fully-connected neural network that classifies handwritten digits (0-9) from the MNIST dataset, built entirely from scratch using PyTorch.

By the end of this project, you'll have:

  • A Python script (train.py) containing your complete training pipeline from data loading to model evaluation.
  • A trained neural network achieving 97%+ accuracy on the MNIST test set.
  • A saved model checkpoint file that can be loaded later for inference or further training.
  • Secret Mission: Replace your fully-connected network with a Convolutional Neural Network and push accuracy above 99%.

Are there any prerequisites?

You should be comfortable with basic Python (variables, functions, loops). You should also understand what a model is and the difference between training and inference.

Set Up Your PyTorch Environment

Before you can train a neural network, you need the right tools on your machine. PyTorch is the deep learning framework that handles everything from tensor math to automatic differentiation.

In this step, you will create a dedicated project folder, install PyTorch inside an isolated Python environment, and confirm everything works with a quick test script.

In this step, get ready to:
  • Create a project folder and Python virtual environment.
  • Install PyTorch and torchvision.
  • Verify the installation by running a test script.
Create a project directory and virtual environment

A virtual environment keeps this project's packages separate from anything else on your machine. This prevents version conflicts between projects.

  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type Visual Studio Code and press Enter to open it.
  • Open the integrated terminal by pressing Ctrl+` (macOS) or Ctrl+` (Windows).
  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir mnist-classifier
cd mnist-classifier
  • Now open this folder in VS Code by clicking File in the top menu bar.
  • Click Open Folder.
  • Navigate to your Desktop, select the mnist-classifier folder, and click Open.

Why open the folder this way?

Opening a folder in VS Code sets it as your workspace root. Any terminal you open afterwards starts inside that folder automatically.

  • Open a new integrated terminal by pressing Ctrl+` (macOS) or Ctrl+` (Windows).
  • Create a virtual environment by running this command:
python3 -m venv venv
  • Activate the virtual environment:

macOS

source venv/bin/activate

Windows

venv\Scripts\activate

You should see (venv) appear at the start of your terminal prompt. This confirms the virtual environment is active.

Don't see (venv) in your prompt?

Make sure you ran the activate command for your operating system. On macOS, use source venv/bin/activate. On Windows, use venv\Scripts\activate.

If you see a permissions error on Windows, try running Set-ExecutionPolicy -Scope CurrentUser -ExecutionPolicy RemoteSigned first, then retry the activate command.

help me activate my Python virtual environment

Install PyTorch and torchvision

PyTorch provides the tensor operations and automatic differentiation engine. torchvision adds common datasets and image transforms you will use to load training data.

  • Install both packages by running this command inside your activated virtual environment:
pip3 install torch torchvision

What does this command do?

pip3 install torch torchvision downloads PyTorch and torchvision from the Python Package Index and installs them into your virtual environment. Because you activated venv first, these packages stay isolated from your system Python.

The download may take a minute or two depending on your internet speed. You should see progress bars followed by a Successfully installed message once it finishes.

Seeing an error during installation?

Make sure (venv) shows in your terminal prompt. If it does not, re-run the activate command from the previous substep before trying the install again.

If you see a network timeout, check your internet connection and retry the same command.

help me install PyTorch with pip

Verify the installation

A quick test script confirms that both packages imported successfully and that PyTorch can create tensors on your machine.

  • Create a new file called train.py by clicking the new file icon in the VS Code Explorer sidebar, then type train.py and press Enter.
  • Add the following verification code to train.py:
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

print(torch.__version__)
print(torch.rand(3, 3))

What does this code do?

The first five lines import PyTorch and torchvision. If either package is missing or broken, Python raises an ImportError immediately.

torch.__version__ prints the installed PyTorch version number so you can confirm which release you have. torch.rand(3, 3) creates a 3x3 tensor filled with random numbers, proving the core tensor engine works.

  • Save train.py by pressing Cmd+S (macOS) or Ctrl+S (Windows).
  • Run the script in your terminal:
python3 train.py

You should see the PyTorch version number on the first line, followed by a 3x3 grid of random decimal numbers.

Seeing ModuleNotFoundError?

This means the package was not found in your current environment. Check that (venv) appears in your terminal prompt. If it does not, activate the virtual environment again and re-run pip3 install torch torchvision.

help me fix ModuleNotFoundError for torch

✔️ Awesome, I've got everything!

Great. Double check you have saved train.py and you are ready to move on.

ⓧ I'd like to double check the full code

Here is the complete train.py at this point:

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

print(torch.__version__)
print(torch.rand(3, 3))

Your environment is ready and PyTorch is confirmed working. Next up, you will load the MNIST handwritten digits dataset and explore what the training data looks like.

Load and Explore the MNIST Dataset

Your PyTorch environment is confirmed and working. Now you need data to train your model on.

The MNIST dataset contains 70,000 images of handwritten digits (0-9). Each image is a 28x28 pixel grayscale picture. In this step, you will load this data, preprocess it into tensors, and organize it into batches for efficient training.

In this step, get ready to:
  • Define a preprocessing pipeline that converts images to tensors and normalizes pixel values.
  • Download the MNIST training and test datasets.
  • Create DataLoaders that feed your model batches of 64 images at a time.
Set up data transforms and download MNIST

Neural networks work with numbers, not raw image files. Before training, each image needs to be converted into a tensor and normalized so pixel values are centered around zero. This helps the model learn faster and more reliably.

  • In train.py, find the two verification lines from the previous step:
print(torch.__version__)
print(torch.rand(3, 3))
  • Delete both of those lines.
  • In their place, add the following code that defines your preprocessing pipeline and downloads the MNIST dataset:
# Data transforms
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

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

What does this code do?

transforms.Compose chains multiple preprocessing steps together into a single pipeline. ToTensor() converts each image from a PIL image (pixel values 0-255) into a PyTorch tensor (values 0.0-1.0).

Normalize((0.1307,), (0.3081,)) shifts and scales the pixel values so they have a mean of 0 and a standard deviation of 1. These specific numbers (0.1307 and 0.3081) are the pre-computed mean and standard deviation of the entire MNIST training set.

datasets.MNIST downloads the dataset into a ./data folder and applies your transform to every image. The train=True flag loads the 60,000 training images, while train=False loads the 10,000 test images.

  • Save train.py.
  • Run the script to trigger the MNIST download by running:
python3 train.py

You should see download progress messages as MNIST downloads into a ./data folder. The script finishes without errors or printed output once the download completes.

Seeing a network error?

MNIST downloads from a remote server. If you see a connection error, check that you are connected to the internet. The download is about 12 MB total.

If you see an ImportError, make sure your virtual environment is activated (you should see (venv) at the start of your terminal prompt).

Help me debug a download error when loading MNIST in PyTorch.

Create DataLoaders for batching

Training on all 60,000 images at once would overwhelm your computer's memory. DataLoaders handle the logistics of feeding data to your model in small, manageable chunks called batches.

  • Add the following code at the bottom of train.py, below the dataset loading lines:
# Create DataLoaders
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

# Verify data loading
train_features, train_labels = next(iter(train_loader))
print(f"Feature batch shape: {train_features.size()}")
print(f"Labels batch shape: {train_labels.size()}")

What does this code do?

DataLoader wraps your dataset and serves up 64 images at a time (one batch) during training. The shuffle=True flag randomizes the order each epoch so the model does not memorize the sequence of images.

The test loader uses a larger batch size of 1000 because evaluation does not need gradients and can process more images simultaneously. shuffle=False keeps the test set in a consistent order for reproducible results.

next(iter(train_loader)) grabs the first batch so you can inspect its shape and confirm everything is wired up correctly.

  • Save train.py and run the script again:
python3 train.py

You should see this output:

Feature batch shape: torch.Size([64, 1, 28, 28])
Labels batch shape: torch.Size([64])

What do those dimensions mean?

The feature shape [64, 1, 28, 28] reads as: 64 images per batch, 1 color channel (grayscale), 28 pixels tall, 28 pixels wide.

The label shape [64] means there are 64 corresponding labels (one digit 0-9 for each image in the batch).

Not seeing the expected output?

If you see a NameError for train_dataset or DataLoader, check that the dataset loading code from the previous substep is directly above the DataLoader code with no missing lines.

If the numbers in your shape output do not match, double check that batch_size=64 is set on your train_loader.

Help me fix the output from my MNIST DataLoader setup.

✔️ Awesome, I've got everything!

Great. Double check you have saved train.py and move on to the next step.

ⓧ I'd like to double check the full code

Here is the complete train.py file at this point:

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

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

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

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

# Verify data loading
train_features, train_labels = next(iter(train_loader))
print(f"Feature batch shape: {train_features.size()}")
print(f"Labels batch shape: {train_labels.size()}")

Your data pipeline is working. MNIST is downloaded, preprocessed, and flowing through DataLoaders in clean batches. Next up, you will define the neural network architecture that learns to recognize these digits.

Define Your Neural Network

Your data pipeline is ready. You have MNIST images flowing through DataLoaders in neat batches of 64.

Now you need something to actually learn from that data. In this step, you will build a neural network class in PyTorch that takes a 28x28 pixel image and predicts which digit (0-9) it represents.

In this step, get ready to:
  • Define a neural network class with three fully-connected layers.
  • Set up a loss function and optimizer to guide training.
Create a model class that subclasses nn.Module

Every neural network in PyTorch is a Python class that inherits from nn.Module. You define your layers in __init__ and describe how data flows through them in a forward method.

  • In train.py, scroll to the bottom of the file (below the batch shape print statements).
  • Add the following class definition:

# Define the model
class MNISTClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(128, 64)
        self.relu2 = nn.ReLU()
        self.fc3 = nn.Linear(64, 10)

    def forward(self, x):
        x = self.flatten(x)
        x = self.relu1(self.fc1(x))
        x = self.relu2(self.fc2(x))
        x = self.fc3(x)
        return x

What does this code do?

This class defines the architecture of your neural network. The __init__ method creates each layer as a named attribute.

  • nn.Flatten() converts the 28x28 image into a flat vector of 784 values.
  • nn.Linear(28 * 28, 128) is the first fully-connected layer. It takes 784 inputs and produces 128 outputs.
  • nn.ReLU() is an activation function that introduces non-linearity, letting the network learn complex patterns.
  • The final layer outputs 10 values. One for each digit (0-9).

The forward method defines the order data passes through these layers. PyTorch calls this automatically when you pass data into the model.

What is nn.Module?

nn.Module is PyTorch's base class for all neural networks. By subclassing it, your model automatically gains the ability to track parameters, move between CPU and GPU, save and load weights, and switch between training and evaluation modes.

  • Save train.py.

You won't see a change yet. The next chunk instantiates the model and prints its structure.

Instantiate the model, loss function, and optimizer

A class definition is just a blueprint. You need to create an instance of it, along with two essential training components: a loss function that measures how wrong the model's predictions are, and an optimizer that updates the model's weights to reduce that loss.

  • Below the class you just added, paste the following code:

# Instantiate model, loss, and optimizer
model = MNISTClassifier()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

print(model)

What does this code do?

  • model = MNISTClassifier() creates a live instance of your network with randomly initialized weights.
  • nn.CrossEntropyLoss() is the standard loss function for classification tasks. It measures how far the model's predictions are from the correct labels.
  • optim.Adam(model.parameters(), lr=0.001) creates an Adam optimizer that will adjust all the model's learnable parameters at a learning rate of 0.001.
  • print(model) displays the full layer structure so you can confirm the architecture is correct.

Why Adam instead of SGD?

Adam adapts its learning rate for each parameter individually, which means it converges faster with default settings. SGD with momentum also works but requires more tuning to get right. For a project running on CPU where training speed matters, Adam at lr=0.001 reliably reaches 97%+ accuracy in just 5 epochs.

  • Save train.py.
  • Run your script by entering this command in your terminal:
python3 train.py

You should see the batch shape prints from earlier, followed by the model summary showing your layers.

Not seeing the model summary?

  • Make sure you pasted the instantiation code OUTSIDE and BELOW the class definition. If it is indented inside the class, Python treats it as a method.
  • Check that your class name matches exactly: MNISTClassifier with that exact capitalization.
  • Verify your virtual environment is still active. You should see (venv) at the start of your terminal prompt.

Help me debug my model definition

✔️ Awesome, I've got everything!

Great. Double check you have saved your file.

ⓧ I'd like to double check the full code

Here is the complete train.py file up to this point:

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

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

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

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

# Verify data loading
train_features, train_labels = next(iter(train_loader))
print(f"Feature batch shape: {train_features.size()}")
print(f"Labels batch shape: {train_labels.size()}")


# Define the model
class MNISTClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(128, 64)
        self.relu2 = nn.ReLU()
        self.fc3 = nn.Linear(64, 10)

    def forward(self, x):
        x = self.flatten(x)
        x = self.relu1(self.fc1(x))
        x = self.relu2(self.fc2(x))
        x = self.fc3(x)
        return x


# Instantiate model, loss, and optimizer
model = MNISTClassifier()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

print(model)

Your neural network is defined and ready to learn. Next up, you will write the training loop that teaches it to recognize digits.

Train the Model

You have a model architecture, a loss function, and an optimizer. But your model's weights are still random. It has never seen a single handwritten digit.

In this step, you will write the training loop that teaches your model to recognize digits. Each pass through the data adjusts the weights a tiny bit, and after five passes (epochs) your network will achieve over 97% accuracy.

In this step, get ready to:
  • Write a training function that performs the forward pass, computes loss, and updates weights.
  • Write an evaluation function that measures accuracy on the test set.
  • Run training for 5 epochs and confirm over 97% test accuracy.
Write the training loop function

The core of every PyTorch project is a four-step cycle that repeats for every batch of data: zero the gradients, run the forward pass, compute the loss, and run backpropagation to update the weights. This pattern is the same whether you are training a tiny MNIST classifier or a billion-parameter language model.

  • In train.py, scroll to the bottom of the file (below the print(model) line).
  • Add the training function by pasting the following code:
# Training function
def train_one_epoch(model, loader, criterion, optimizer):
    model.train()
    running_loss = 0.0
    for data, target in loader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * data.size(0)
    return running_loss / len(loader.dataset)

What does this code do?

This function encapsulates the full training loop pattern that every PyTorch project uses.

  • model.train() sets the model to training mode. This ensures layers like Dropout and BatchNorm behave correctly during training.
  • optimizer.zero_grad() resets all gradients to zero. Without this, gradients would accumulate from the previous batch.
  • output = model(data) is the forward pass. Data flows through each layer of your network.
  • loss.backward() computes how much each weight contributed to the error. This is backpropagation.
  • optimizer.step() nudges each weight in the direction that reduces the loss.
  • Save train.py.
  • You won't see a change yet. The training loop below calls this function to actually train the model.

Seeing an IndentationError?

Make sure the entire function body is indented with four spaces under def train_one_epoch. The for loop body should be indented another four spaces (eight total).

Also check that there is a blank line between print(model) and the comment # Training function.

Why is my training function not working?

Write the evaluation function and run training

Training adjusts the weights, but how do you know the model is actually learning? You need a separate function that tests the model on data it has never trained on. This gives an honest measure of how well the network generalizes.

  • In train.py, add the evaluation function directly below the train_one_epoch function:
# Evaluation function
def evaluate(model, loader):
    model.eval()
    correct = 0
    with torch.no_grad():
        for data, target in loader:
            output = model(data)
            pred = output.argmax(dim=1)
            correct += pred.eq(target).sum().item()
    return 100.0 * correct / len(loader.dataset)

What does this code do?

  • model.eval() switches the model to evaluation mode. This disables training-specific behaviors like Dropout.
  • torch.no_grad() disables gradient computation entirely. This saves memory and speeds up inference since you do not need gradients when evaluating.
  • output.argmax(dim=1) picks the digit (0-9) with the highest confidence score for each image in the batch.
  • The function returns accuracy as a percentage of correctly classified test images.
  • Save train.py.
  • You won't see a change yet. The training loop below calls both functions together.

Confused about model.train() vs model.eval()?

These two calls control different phases. model.train() enables training behavior (active Dropout, updating BatchNorm stats). model.eval() disables those behaviors so you get consistent, deterministic predictions during testing.

Help me understand when to use model.train() vs model.eval() in PyTorch.

Now you will wire everything together. The training loop calls train_one_epoch to update the weights and evaluate to measure progress after each epoch.

  • Add the training loop below the evaluate function:
# Train for 5 epochs
for epoch in range(1, 6):
    loss = train_one_epoch(model, train_loader, criterion, optimizer)
    accuracy = evaluate(model, test_loader)
    print(f"Epoch {epoch}: Loss = {loss:.4f}, Test Accuracy = {accuracy:.2f}%")

What does this code do?

This loop trains the model for 5 complete passes through the training data. Each pass is called an epoch. After each epoch, it evaluates the model on the test set and prints both the average loss and the test accuracy.

You should see the loss decrease and the accuracy increase with each epoch. By epoch 5, the model should exceed 97% accuracy on digits it has never seen before.

  • Save train.py.
  • Run the script by executing:
python3 train.py
  • You should see the batch shape prints and model summary from earlier, followed by 5 lines of training output. The final line should show test accuracy above 97%. Training takes approximately 30-60 seconds on most machines.

What should I see?

Your output should look similar to this (exact numbers will vary slightly due to random initialization):

Epoch 1: Loss = 0.2345, Test Accuracy = 96.12%
Epoch 2: Loss = 0.0987, Test Accuracy = 97.05%
Epoch 3: Loss = 0.0654, Test Accuracy = 97.45%
Epoch 4: Loss = 0.0478, Test Accuracy = 97.62%
Epoch 5: Loss = 0.0356, Test Accuracy = 97.80%

The loss should decrease each epoch and accuracy should climb above 97% by the final epoch.

Accuracy stuck below 90%?

Check that your train_one_epoch function includes optimizer.zero_grad() before the forward pass. Without it, gradients accumulate across batches and training becomes unstable.

Also verify that loss.backward() comes before optimizer.step(). The optimizer needs computed gradients before it can update the weights.

My model accuracy is not improving during training. Help me debug.

✔️ Awesome, I've got everything!

Great. Double check that you have saved train.py and your final epoch shows accuracy above 97%.

ⓧ I'd like to double check the full code

Here is the full train.py file at this point. Compare it with yours:

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

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

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

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

# Verify data loading
train_features, train_labels = next(iter(train_loader))
print(f"Feature batch shape: {train_features.size()}")
print(f"Labels batch shape: {train_labels.size()}")


# Define the model
class MNISTClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(128, 64)
        self.relu2 = nn.ReLU()
        self.fc3 = nn.Linear(64, 10)

    def forward(self, x):
        x = self.flatten(x)
        x = self.relu1(self.fc1(x))
        x = self.relu2(self.fc2(x))
        x = self.fc3(x)
        return x


# Instantiate model, loss, and optimizer
model = MNISTClassifier()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

print(model)


# Training function
def train_one_epoch(model, loader, criterion, optimizer):
    model.train()
    running_loss = 0.0
    for data, target in loader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * data.size(0)
    return running_loss / len(loader.dataset)


# Evaluation function
def evaluate(model, loader):
    model.eval()
    correct = 0
    with torch.no_grad():
        for data, target in loader:
            output = model(data)
            pred = output.argmax(dim=1)
            correct += pred.eq(target).sum().item()
    return 100.0 * correct / len(loader.dataset)


# Train for 5 epochs
for epoch in range(1, 6):
    loss = train_one_epoch(model, train_loader, criterion, optimizer)
    accuracy = evaluate(model, test_loader)
    print(f"Epoch {epoch}: Loss = {loss:.4f}, Test Accuracy = {accuracy:.2f}%")

Your network just went from random guesses to 97%+ accuracy in under a minute. Next up, you will save this trained model to disk and load it back to run predictions on individual digits.

Save the Model and Run Inference

Your model has been trained and lives in memory. If you close the script or restart your computer, all that training work disappears.

In this step, you'll save your trained model's weights to disk and then load them into a fresh model instance to prove inference works independently of training.

In this step, get ready to:
  • Save your trained model's weights to a .pth file.
  • Load the saved weights into a new model instance and run a prediction on a single digit.
Save the trained model

PyTorch separates a model's architecture (the class definition) from its learned weights (the numbers computed during training). The state_dict() method returns a dictionary containing all of those learned weights.

  • In train.py, find the training loop at the bottom of your file:
# Train for 5 epochs
for epoch in range(1, 6):
    loss = train_one_epoch(model, train_loader, criterion, optimizer)
    accuracy = evaluate(model, test_loader)
    print(f"Epoch {epoch}: Loss = {loss:.4f}, Test Accuracy = {accuracy:.2f}%")
  • Add the following code directly below the training loop to save the model weights:
# Save the model
torch.save(model.state_dict(), 'mnist_model.pth')
print("Model saved to mnist_model.pth")

What does this code do?

model.state_dict() extracts all the learned weights and biases from your trained model as a Python dictionary.

torch.save() serializes that dictionary to a file called mnist_model.pth. This file contains only the weights, not the architecture. To use the model later, you need both the class definition and this file.

You won't see a change yet. The next chunk loads this file back and runs a prediction.

Load the model and run inference on a single sample

In production, you train a model once and then deploy it separately for inference. Loading weights into a fresh model instance simulates that exact workflow.

  • Add the following code below the save lines you just added:
# Load and run inference
loaded_model = MNISTClassifier()
loaded_model.load_state_dict(torch.load('mnist_model.pth', weights_only=True))
loaded_model.eval()

# Grab a single test sample
sample_image, sample_label = test_dataset[0]
with torch.no_grad():
    prediction = loaded_model(sample_image.unsqueeze(0))
    predicted_digit = prediction.argmax(dim=1).item()

print(f"Predicted: {predicted_digit}, Actual: {sample_label}")

What does this code do?

  • MNISTClassifier() creates a brand-new model with random weights. load_state_dict() then overwrites those random weights with the trained weights from your saved file.
  • weights_only=True is a security parameter that prevents arbitrary code execution when loading model files from untrusted sources.
  • sample_image.unsqueeze(0) adds a batch dimension. The model expects input shaped [batch_size, 1, 28, 28], but a single image is shaped [1, 28, 28]. Unsqueeze wraps it in a batch of one.
  • prediction.argmax(dim=1).item() finds which of the 10 output neurons has the highest value, giving you the predicted digit as a plain Python integer.
  • Save train.py.
  • Run the script by entering this command in your terminal:
python3 train.py
  • After the training completes (about 30-60 seconds), you should see Model saved to mnist_model.pth followed by a line like Predicted: 7, Actual: 7 (your digit may differ, but predicted and actual should match).

Seeing a mismatch between predicted and actual?

A rare mismatch on a single sample is normal. Your model is 97%+ accurate, not 100%. Run the script again and check that the overall test accuracy is above 97%. If accuracy is very low (below 90%), make sure your training loop runs for 5 epochs and that the optimizer is optim.Adam with lr=0.001.

Help me debug my model's low accuracy

  • Confirm the model file exists by running:
ls mnist_model.pth
  • You should see mnist_model.pth in the output, confirming your trained weights are saved to disk.

✔️ Awesome, I've got everything!

Great. Double check you've saved your file.

ⓧ I'd like to double check the full code

Here is the complete train.py file:

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

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

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

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

# Verify data loading
train_features, train_labels = next(iter(train_loader))
print(f"Feature batch shape: {train_features.size()}")
print(f"Labels batch shape: {train_labels.size()}")


# Define the model
class MNISTClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(128, 64)
        self.relu2 = nn.ReLU()
        self.fc3 = nn.Linear(64, 10)

    def forward(self, x):
        x = self.flatten(x)
        x = self.relu1(self.fc1(x))
        x = self.relu2(self.fc2(x))
        x = self.fc3(x)
        return x


# Instantiate model, loss, and optimizer
model = MNISTClassifier()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

print(model)


# Training function
def train_one_epoch(model, loader, criterion, optimizer):
    model.train()
    running_loss = 0.0
    for data, target in loader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * data.size(0)
    return running_loss / len(loader.dataset)


# Evaluation function
def evaluate(model, loader):
    model.eval()
    correct = 0
    with torch.no_grad():
        for data, target in loader:
            output = model(data)
            pred = output.argmax(dim=1)
            correct += pred.eq(target).sum().item()
    return 100.0 * correct / len(loader.dataset)


# Train for 5 epochs
for epoch in range(1, 6):
    loss = train_one_epoch(model, train_loader, criterion, optimizer)
    accuracy = evaluate(model, test_loader)
    print(f"Epoch {epoch}: Loss = {loss:.4f}, Test Accuracy = {accuracy:.2f}%")

# Save the model
torch.save(model.state_dict(), 'mnist_model.pth')
print("Model saved to mnist_model.pth")

# Load and run inference
loaded_model = MNISTClassifier()
loaded_model.load_state_dict(torch.load('mnist_model.pth', weights_only=True))
loaded_model.eval()

# Grab a single test sample
sample_image, sample_label = test_dataset[0]
with torch.no_grad():
    prediction = loaded_model(sample_image.unsqueeze(0))
    predicted_digit = prediction.argmax(dim=1).item()

print(f"Predicted: {predicted_digit}, Actual: {sample_label}")

Why save only the state_dict and not the whole model?

Saving the full model with torch.save(model, path) uses Python's pickle under the hood, which ties the saved file to your exact class definitions and directory structure. If you rename a file or refactor your code, the model breaks.

Saving just the state_dict() gives you a portable dictionary of weights. You can load it into any compatible architecture, even in a completely different project.

Secret mission

Build a CNN for 99%+ Accuracy

Your fully-connected network hits 97% accuracy. Can you push past 99% by leveraging the spatial structure of images? Replace your flat network with a Convolutional Neural Network that sees patterns the way your eyes do.

Clean Up Your Resources

Clean Up Your Resources

Decide whether to keep your resources running, pause them to come back later, or delete them entirely. This project runs entirely locally, so there are no ongoing costs.

Resources you used:

  • mnist-classifier project folder (contains train.py, model files, downloaded MNIST data, and virtual environment).

Keep everything running

No action needed. Choose this if you plan to keep experimenting with your model or want to tweak the architecture.

  • Your mnist-classifier folder stays on your Desktop with all files intact.
  • The downloaded MNIST dataset (~60 MB in the data/ folder) stays cached so future runs skip the download.
  • Your virtual environment keeps PyTorch and torchvision installed without affecting your system Python.

Pause - I'll come back to this later

Deactivate your virtual environment to free up your terminal session. All files remain on disk for when you return.

  • Deactivate the virtual environment by running:
deactivate
  • When you come back, navigate to your project folder and reactivate with source venv/bin/activate.

Delete - I don't want to use this again

Remove all project files, the virtual environment, downloaded data, and trained models.

  • Deactivate the virtual environment first (if still active):
deactivate
  • Delete the entire project folder by running:

macOS

rm -rf ~/Desktop/mnist-classifier

Windows

rmdir /s /q %USERPROFILE%\Desktop\mnist-classifier

This removes train.py, both .pth model files, the downloaded MNIST dataset, and the virtual environment with all installed packages.

Nice Work!

Nice Work!

You just trained a neural network from scratch that recognizes handwritten digits with over 97% accuracy. That is real deep learning, built entirely by you.

You've learned how to:

  • Trained a neural network from scratch in PyTorch that classifies handwritten digits with 97%+ accuracy.
  • Mastered the core PyTorch training loop pattern: forward pass, compute loss, backward pass, optimizer step.
  • Saved and loaded a model checkpoint for inference, proving your trained model can be reused without retraining.
  • Secret Mission: Built a CNN that leverages spatial features to achieve 99%+ accuracy, demonstrating why convolutional architectures dominate image tasks.

Ready to quiz yourself?