Train a Neural Network with PyTorch

Train a neural network from scratch to classify clothing images.

Introduction

30 Second Summary

You can glance at a pile of laundry and instantly tell a sneaker from a sandal. Teaching a computer to do the same thing requires breaking that intuition down into math.

In this project, you will train a neural network from scratch using PyTorch to classify images of clothing into 10 categories. You will watch it improve from random guessing to ~85% accuracy in just a few minutes of training on your CPU.

What You'll Build

A neural network that takes 28x28 grayscale images of clothing (t-shirts, sneakers, dresses) and predicts which of 10 garment types they belong to.

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

  • A trained neural network that classifies clothing images at ~85% accuracy.
  • A complete Jupyter notebook showing the full ML workflow from data loading to prediction visualization.
  • Hands-on understanding of how PyTorch tensors, layers, loss functions, and optimizers work together to train a model.
  • Secret Mission: Add a confusion matrix and build a deeper network with dropout regularization to improve model performance.

Are there any prerequisites?

You should be comfortable with basic Python (functions, loops, classes) and familiar with using the terminal. No prior machine learning experience is needed.

Before We Start

Before We Start...

Set Up Python with pyenv

Training a neural network requires a modern version of Python. Your system currently has Python 3.4, which is far too old for PyTorch and the other ML libraries you will use.

To get the right version without breaking your system Python, you will use pyenv. It lets you install and switch between Python versions per project.

In this step, get ready to:
  • Install pyenv and its build dependencies.
  • Install Python 3.12 and set it as this project's version.
  • Create a virtual environment and install the Jupyter extension in VS Code.
Install pyenv build dependencies and pyenv

pyenv compiles Python from source, so it needs a set of system libraries installed first.

  • Open a terminal in VS Code by pressing Ctrl+`.
  • Install the required build dependencies by running this command:
sudo apt install -y build-essential zlib1g-dev libssl-dev libreadline-dev liblzma-dev libsqlite3-dev libffi-dev libbz2-dev

What does this command do?

These packages provide the C libraries and headers that Python needs when compiling from source. Without them, the build would fail with missing-header errors.

  • Install pyenv using the official installer by running this command:
curl -fsSL https://pyenv.run | bash

What does this command do?

This downloads and runs the pyenv installer script. It clones pyenv into ~/.pyenv and sets up useful plugins like pyenv-virtualenv.

After the installer finishes, you need to configure your shell so it can find pyenv.

  • Add pyenv to your shell startup by running these commands:
echo 'export PYENV_ROOT="$HOME/.pyenv"' >> ~/.bashrc
echo '[[ -d $PYENV_ROOT/bin ]] && export PATH="$PYENV_ROOT/bin:$PATH"' >> ~/.bashrc
echo 'eval "$(pyenv init - bash)"' >> ~/.bashrc

What do these lines do?

They tell your shell where pyenv lives and activate its shim system. The shims intercept calls to python3 and route them to whichever version pyenv has selected for the current directory.

  • Reload your shell configuration by running:
exec "$SHELL"
  • Verify pyenv is working by running:
pyenv --version

✔️ I see a version number

pyenv is installed and ready. You should see something like pyenv 2.6.27 or similar. Move on to the next substep.

ⓧ Command not found

Make sure you ran all three echo commands above and that each completed without errors.

  • Run exec "$SHELL" again to reload your shell.
  • If the problem persists, check that ~/.pyenv exists by running ls ~/.pyenv. If the folder is missing, re-run the curl installer command.

help me fix pyenv not being found after installation

Install Python 3.12 and set it as the local version

With pyenv ready, you can now compile and install a modern Python without touching your system installation.

  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir fashion-classifier
cd fashion-classifier
  • Confirm the folder exists by running:
ls ~/Desktop

You should see fashion-classifier in the output.

  • Install Python 3.12 by running this command:
pyenv install 3.12

What does this command do?

pyenv uses prefix auto-resolution. Giving it 3.12 tells it to download and compile the latest patch version in the 3.12 line (for example, 3.12.9). This takes a few minutes because Python is being built from source on your machine.

This will take a few minutes

You will see download and compilation output scrolling by. Wait until you see a line confirming the install is complete before continuing.

  • Set Python 3.12 as the version for this project folder by running:
pyenv local 3.12

What does pyenv local do?

This writes the version number to a .python-version file in the current directory. Whenever you are inside fashion-classifier, pyenv automatically uses that Python version instead of the system default.

  • Verify the correct version is active by running:
python3 --version

✔️ I see Python 3.12.x

Python 3.12 is active for this project. The exact patch number (for example, 3.12.9) depends on the latest release when you ran pyenv install. Move on to the next substep.

ⓧ I see an older version

Your shell may still be pointing at the system Python. This usually means pyenv's shims are not active in your current session.

  • Confirm you are inside the fashion-classifier directory by running pwd. The .python-version file only applies when your terminal is in that folder.
  • Run exec "$SHELL" to reload your shell, then try python3 --version again.

help me get pyenv local to work correctly

ⓧ Command not found

If python3 is not found at all, pyenv may not have finished installing correctly.

  • Run pyenv versions to check whether any Python version is listed.
  • If no version appears, re-run pyenv install 3.12 and watch for any error messages during compilation.
  • After the install succeeds, run pyenv local 3.12 again and check with python3 --version.

help me troubleshoot pyenv install failing

Create a virtual environment and install the Jupyter extension

A virtual environment isolates this project's packages from the rest of your system. This prevents version conflicts between projects.

  • Create a virtual environment by running this command:
python3 -m venv venv
  • Activate the virtual environment by running:
source venv/bin/activate

Your terminal prompt should now show (venv) at the beginning. This confirms you are working inside the virtual environment.

What is a virtual environment?

A virtual environment is a self-contained directory that holds its own Python binary and packages. When activated, pip install puts libraries inside this folder instead of globally.

This keeps your project's dependencies separate. If you work on another project that needs different library versions, they will not collide.

Not seeing (venv) in your prompt?

Make sure the source command completed without errors. If you see a file-not-found message, confirm you are still inside the fashion-classifier directory by running pwd.

help me activate my Python virtual environment

  • Now install the Jupyter extension in VS Code. Press Ctrl+Shift+X to open the Extensions sidebar.
  • Type Jupyter into the search bar.
  • Click Install on the extension named Jupyter published by Microsoft.

Why do I need the Jupyter extension?

Jupyter notebooks let you run Python code in small cells and see results immediately. This is the standard way to experiment with machine learning models because you can re-run training steps, tweak parameters, and visualize outputs without restarting from scratch.

Can't find the Jupyter extension?

Make sure you are searching in the Extensions sidebar (the square-blocks icon on the left), not the file explorer. The extension publisher should say Microsoft.

help me find and install the Jupyter extension in VS Code

Your environment is ready. Next up, you will install PyTorch and create the notebook where your neural network will live.

Install PyTorch and Create Your Notebook

Your Python environment is ready and waiting. Now it's time to fill it with the machine learning tools you need and start building your notebook.

In this step, you'll install PyTorch (the deep learning framework), load the Fashion-MNIST dataset, and explore what 70,000 images of clothing look like as data.

In this step, get ready to:
  • Install PyTorch, torchvision, and matplotlib in your virtual environment.
  • Create a Jupyter notebook and verify PyTorch is working.
  • Load and explore the Fashion-MNIST dataset.
Install PyTorch, torchvision, and matplotlib

PyTorch is the framework that handles all the heavy lifting for building and training neural networks. You also need torchvision (for the dataset) and matplotlib (for visualizations).

  • Install PyTorch and torchvision by running this command in your terminal:
pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cpu

Why the --index-url flag?

PyTorch distributes CPU-only and GPU versions from different package indexes. The --index-url flag points pip to the CPU-only packages, which are smaller and faster to download.

  • Install matplotlib by running:
pip install matplotlib

You should see a success message for each package once the installations finish.

Seeing an installation error?

Make sure your virtual environment is active. You should see (venv) at the start of your terminal prompt. If not, activate it with source venv/bin/activate and try again.

Help me troubleshoot my PyTorch installation

Create a Jupyter notebook and verify PyTorch

A Jupyter notebook lets you write and run code in small chunks called cells. This is perfect for machine learning because you can run each step independently and see results immediately.

  • Open the Command Palette by pressing Cmd+Shift+P (macOS) or Ctrl+Shift+P (Windows/Linux).
  • Type Create: New Jupyter Notebook and press Enter.
  • A new untitled notebook opens. Save it immediately by pressing Cmd+S (macOS) or Ctrl+S (Windows/Linux).
  • Name the file fashion_mnist_classifier.ipynb and save it in your project folder.

What about the kernel?

VS Code will ask you to select a kernel (the Python environment that runs your code). Click Select Kernel in the top right of the notebook, choose Python Environments, and select the venv environment you created in Step 1.

  • In the first cell of your notebook, paste the following import code:
# Cell 1: Imports
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt

What does this code do?

  • torch is the core PyTorch library for tensor operations and automatic differentiation.
  • torch.nn provides building blocks for neural networks (layers, activation functions).
  • torch.optim contains optimizers that update the network's weights during training.
  • DataLoader handles batching and shuffling your data for training.
  • torchvision.datasets gives you access to common image datasets like Fashion-MNIST.
  • matplotlib.pyplot is for creating charts and displaying images.
  • Run the cell by pressing Shift+Enter.

If the cell runs without errors, PyTorch is installed correctly and your notebook is ready to go.

Seeing ModuleNotFoundError?

This means the notebook is using a different Python environment than the one where you installed PyTorch. Click Select Kernel in the top right corner and make sure you've selected the venv environment from your project folder.

Help me fix my notebook kernel selection

Load the Fashion-MNIST dataset and explore its structure

Fashion-MNIST is a dataset of 70,000 grayscale images of clothing items. Each image is 28x28 pixels and belongs to one of 10 categories (like t-shirts, sneakers, and dresses).

  • Create a new cell below Cell 1 by clicking the + Code button, then paste this code:
# Cell 2: Load Fashion-MNIST dataset
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)

class_names = [
    "T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",
    "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"
]

print(f"Training set size: {len(training_data)}")
print(f"Test set size: {len(test_data)}")
print(f"Image shape: {training_data[0][0].shape}")
print(f"Number of classes: {len(class_names)}")

What does this code do?

  • datasets.FashionMNIST downloads the dataset into a data/ folder in your project.
  • train=True loads the 60,000 training images. train=False loads the 10,000 test images.
  • transform=ToTensor() converts each image from a PIL image to a PyTorch tensor (numbers the network can process).
  • class_names maps the numeric labels (0-9) to human-readable clothing categories.
  • Run Cell 2 by pressing Shift+Enter.

The first run downloads the dataset (about 30 MB). You should see output showing 60,000 training images, 10,000 test images, image shape torch.Size([1, 28, 28]), and 10 classes.

Download taking too long?

The download is about 30 MB total. If your connection is slow, just let it finish. Once downloaded, the data is cached in the data/ folder and won't need to be downloaded again.

Help me troubleshoot the Fashion-MNIST download

  • Create a new cell and paste this visualization code:
# Cell 3: Visualize sample images
fig, axes = plt.subplots(2, 5, figsize=(12, 5))
for i, ax in enumerate(axes.flat):
    image, label = training_data[i]
    ax.imshow(image.squeeze(), cmap="gray")
    ax.set_title(class_names[label])
    ax.axis("off")
plt.tight_layout()
plt.show()

What does this code do?

This creates a 2x5 grid showing the first 10 images from the training set. Each image is displayed in grayscale with its class name as the title. The .squeeze() removes the extra channel dimension so matplotlib can display it as a 2D image.

  • Run Cell 3 by pressing Shift+Enter.

You should see a grid of 10 grayscale images showing different clothing items with labels like "Ankle boot", "T-shirt/top", and "Trouser" above each one.

  • Create one more cell and paste this code to set up the DataLoaders:
# Cell 4: Create DataLoaders
batch_size = 64

train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=False)

What are DataLoaders?

A DataLoader wraps your dataset and feeds it to the neural network in small batches (64 images at a time). Training on batches is more memory-efficient than loading all 60,000 images at once.

shuffle=True randomizes the training data order each epoch so the network doesn't memorize the sequence. The test data stays unshuffled so results are reproducible.

  • Run Cell 4 by pressing Shift+Enter.

This cell runs silently with no output. The DataLoaders are now ready to feed batches into your neural network in the next step.

✔️ Awesome, I've got everything!

Great. Double check you've saved your notebook with Cmd+S (macOS) or Ctrl+S (Windows/Linux).

ⓧ I'd like to double check the full code

Here is the complete notebook code up to this point. Each comment starting with # Cell marks the beginning of a new cell in your notebook.

# Cell 1: Imports
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt

# Cell 2: Load Fashion-MNIST dataset
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)

class_names = [
    "T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",
    "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"
]

print(f"Training set size: {len(training_data)}")
print(f"Test set size: {len(test_data)}")
print(f"Image shape: {training_data[0][0].shape}")
print(f"Number of classes: {len(class_names)}")

# Cell 3: Visualize sample images
fig, axes = plt.subplots(2, 5, figsize=(12, 5))
for i, ax in enumerate(axes.flat):
    image, label = training_data[i]
    ax.imshow(image.squeeze(), cmap="gray")
    ax.set_title(class_names[label])
    ax.axis("off")
plt.tight_layout()
plt.show()

# Cell 4: Create DataLoaders
batch_size = 64

train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=False)

Your data pipeline is ready. You have 60,000 training images flowing through DataLoaders in batches of 64. Next up, you'll build the neural network that learns to classify these images.

Build Your Neural Network

You've loaded and explored 70,000 images of clothing, and your DataLoaders are ready to feed data in batches.

Now it's time to build the brain of this project: a neural network that transforms raw pixel data into clothing predictions. Each image is 28x28 pixels, which means 784 numbers go in, and 10 class probabilities come out.

In this step, get ready to:
  • Define a neural network class in PyTorch.
  • Instantiate the model and inspect its architecture.
  • Configure the loss function and optimizer for training.
Define the neural network class

A neural network is a stack of mathematical layers. Each layer applies a linear transformation and an activation function to its input. By stacking multiple layers, the network can learn increasingly complex patterns in data.

  • In your notebook, create a new code cell below Cell 4 and add the following code:
# Cell 5: Define the neural network
class FashionClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.network = nn.Sequential(
            nn.Linear(28 * 28, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 10),
        )

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

What does this code do?

This defines a Python class called FashionClassifier that inherits from nn.Module. Every PyTorch model is built this way.

  • nn.Flatten() reshapes a 28x28 image into a flat list of 784 numbers.
  • nn.Linear(28 * 28, 512) is the first layer. It takes 784 inputs and produces 512 outputs.
  • nn.ReLU() is an activation function that introduces non-linearity. Without it, stacking layers would be no better than a single layer.
  • The second hidden layer compresses 512 values down to 256.
  • The final layer outputs 10 values, one score per clothing class.
  • The forward() method defines how data flows through the network: flatten first, then pass through all layers.

You won't see any output yet. The next substep instantiates the model and prints its structure.

Instantiate the model and print its architecture

Defining a class creates a blueprint. To actually use the network, you need to create an instance of it.

  • At the bottom of Cell 5 (below the class definition), add these two lines:
model = FashionClassifier()
print(model)
  • Run Cell 5 by clicking the play button to the left of the cell (or pressing Shift+Enter).

What should I see?

The output shows the full architecture of your model. You should see FashionClassifier with a flatten layer and a network block containing the sequence of Linear and ReLU layers with their input and output sizes listed.

This confirms PyTorch registered all of your layers and their learnable parameters.

Seeing a NameError?

Make sure you ran Cell 1 (the imports cell) first. The nn module needs to be imported before you can use nn.Module or nn.Linear.

Also check that your class definition has no indentation errors. Each method inside the class must be indented one level, and the code inside each method indented two levels.

help me fix my FashionClassifier class definition

Set up the loss function and optimizer

The model can now make predictions, but they're random. To teach it, you need two things: a way to measure how wrong it is (the loss function), and a way to adjust its weights to reduce that wrongness (the optimizer).

  • Create a new code cell below Cell 5 and add the following:
# Cell 6: Set up loss function and optimizer
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

What does this code do?

  • CrossEntropyLoss measures the difference between the model's predicted probabilities and the correct answer. It's the standard loss function for classification tasks.
  • optim.SGD is Stochastic Gradient Descent. It updates the model's weights after each batch to reduce the loss.
  • model.parameters() passes all learnable weights from your network to the optimizer so it knows what to update.
  • lr=0.01 is the learning rate. It controls how big each weight adjustment is.
  • momentum=0.9 helps the optimizer move consistently in one direction rather than bouncing around, speeding up convergence.
  • Run Cell 6 by pressing Shift+Enter.

This cell produces no visible output. That's expected. The loss function and optimizer are now configured and ready for the training loop.

Seeing a NameError for optim?

The optim module is imported in Cell 1 as import torch.optim as optim. Make sure you ran Cell 1 before running Cell 6.

help me fix my loss function and optimizer setup

✔️ Awesome, I've got everything!

Great. Double check that you've saved your notebook.

ⓧ I'd like to double check the full code

Here's the complete notebook up to this point:

# Cell 1: Imports
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt

# Cell 2: Load Fashion-MNIST dataset
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)

class_names = [
    "T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",
    "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"
]

print(f"Training set size: {len(training_data)}")
print(f"Test set size: {len(test_data)}")
print(f"Image shape: {training_data[0][0].shape}")
print(f"Number of classes: {len(class_names)}")

# Cell 3: Visualize sample images
fig, axes = plt.subplots(2, 5, figsize=(12, 5))
for i, ax in enumerate(axes.flat):
    image, label = training_data[i]
    ax.imshow(image.squeeze(), cmap="gray")
    ax.set_title(class_names[label])
    ax.axis("off")
plt.tight_layout()
plt.show()

# Cell 4: Create DataLoaders
batch_size = 64

train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=False)

# Cell 5: Define the neural network
class FashionClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.network = nn.Sequential(
            nn.Linear(28 * 28, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 10),
        )

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

model = FashionClassifier()
print(model)

# Cell 6: Set up loss function and optimizer
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

Your neural network is defined, instantiated, and configured with a loss function and optimizer. Next up, you'll write the training loop that teaches this model to recognize clothing.

Train Your Model

Your neural network has a structure, but right now its weights are completely random. It would score about 10% accuracy on classifying clothing images, no better than a coin flip across 10 categories.

Training is the process of feeding the model thousands of images, measuring how wrong its predictions are, and nudging the weights in a better direction. In this step, you will write a training loop, an evaluation function, and run the full training process to watch your model learn.

In this step, get ready to:
  • Write a training loop that computes loss and updates weights via backpropagation.
  • Run training for 10 epochs and watch the model improve from random guessing to ~85% accuracy.
  • Evaluate the trained model on the test set and report overall accuracy.
Write the training loop

A training loop is the core engine of deep learning. Each pass through the loop feeds a batch of images to the model, measures how wrong the predictions are (the loss), then adjusts the weights to reduce that error.

  • Create a new cell below Cell 6 in your notebook (Cell 7) and paste the following code:
# Cell 7: Define training function
def train_one_epoch(dataloader, model, loss_fn, optimizer):
    model.train()
    running_loss = 0.0
    for batch, (X, y) in enumerate(dataloader):
        pred = model(X)
        loss = loss_fn(pred, y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    avg_loss = running_loss / len(dataloader)
    return avg_loss

What does this code do?

This function trains the model for one complete pass through the dataset (one epoch). Here is what each key line does:

  • model.train() puts the model in training mode so it tracks gradients.
  • pred = model(X) feeds a batch of images through the network to get predictions (forward propagation).
  • loss.backward() computes how much each weight contributed to the error (backpropagation).
  • optimizer.step() nudges the weights in the direction that reduces the loss.

What is backpropagation?

Backpropagation is the algorithm that figures out which weights caused the error. It works backwards from the output layer to the input layer, calculating a gradient (a direction to improve) for every single weight in the network.

PyTorch's autograd system handles this automatically when you call loss.backward(). You define the architecture, and PyTorch figures out the math.

  • Run this cell. You should see no output (the function is defined but not called yet).

Seeing a syntax error?

  • Check that the indentation uses consistent spaces (4 spaces per level).
  • Make sure you have a colon after def train_one_epoch(dataloader, model, loss_fn, optimizer):.
  • Help me debug my training function
Set up the evaluation function

To know whether training is working, we need a way to measure accuracy on images the model has never seen. This evaluation function tests the model against the test set after each epoch.

  • Create a new cell (Cell 8) and paste the following code:
# Cell 8: Define evaluation function
def evaluate(dataloader, model, loss_fn):
    model.eval()
    correct = 0
    total = 0
    test_loss = 0.0

    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).sum().item()
            total += y.size(0)

    accuracy = correct / total
    avg_loss = test_loss / len(dataloader)
    return accuracy, avg_loss

What does this code do?

  • model.eval() switches off training-specific behavior so the model gives consistent predictions.
  • torch.no_grad() tells PyTorch not to compute gradients during evaluation. This saves memory and speeds things up since we are only measuring, not learning.
  • pred.argmax(1) == y checks which predictions matched the true labels. The model outputs 10 scores (one per class), and argmax picks the highest-scoring class as the prediction.
  • Run this cell. Again, no output is expected since the function is only defined here.

Getting a NameError for torch?

  • Make sure Cell 1 (your imports cell) has been run successfully before running this cell.
  • If you restarted your kernel, re-run all previous cells from Cell 1 through Cell 6 first.
  • Help me fix my evaluate function
Run training for 10 epochs and evaluate

Now you have both pieces: a function that trains for one epoch, and a function that measures accuracy. Time to put them together and watch your model learn.

  • Create a new cell (Cell 9) and paste the following code to run the full training loop:
# Cell 9: Train the model
epochs = 10

for epoch in range(epochs):
    train_loss = train_one_epoch(train_dataloader, model, loss_fn, optimizer)
    accuracy, test_loss = evaluate(test_dataloader, model, loss_fn)
    print(f"Epoch {epoch+1:>2d} | Train Loss: {train_loss:.4f} | Test Loss: {test_loss:.4f} | Accuracy: {accuracy*100:.1f}%")

print("\nTraining complete!")

What does this code do?

  • epochs = 10 means the model sees the entire 60,000-image training set 10 times.
  • Each iteration calls train_one_epoch() to train, then evaluate() to measure progress on unseen test data.
  • The printed line shows train loss (how well it fits training data), test loss (how well it generalizes), and accuracy (percentage of correct predictions on the test set).
  • Run this cell. Training takes about 1-2 minutes on CPU.

You should see 10 lines printed, one per epoch. Watch the train loss decrease and accuracy increase with each epoch. Your final accuracy should land around 85%.

What should I see?

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

Epoch  1 | Train Loss: 0.6523 | Test Loss: 0.4812 | Accuracy: 82.3%
Epoch  2 | Train Loss: 0.4105 | Test Loss: 0.4201 | Accuracy: 84.1%
...
Epoch 10 | Train Loss: 0.2451 | Test Loss: 0.3512 | Accuracy: 86.2%

Training complete!

The key pattern is that train loss should consistently decrease. Test accuracy should climb from around 80% to roughly 85-88%.

Why does loss decrease but accuracy doesn't always go up?

Loss measures how confident the model is in its answers, not just whether it got them right. Sometimes the model becomes more confident on correct answers (lower loss) without changing any predictions (same accuracy).

Occasional accuracy dips between epochs are normal. The model is still learning. What matters is the overall trend across all 10 epochs.

Cell seems stuck or taking too long?

  • Training on CPU is expected to take 1-2 minutes total. If no output appears after 30 seconds, wait. Output prints after each full epoch, not after each batch.
  • If you see a NameError for train_one_epoch or evaluate, make sure you ran Cell 7 and Cell 8 first.
  • If accuracy stays near 10% and never improves, your model or optimizer may not have been set up correctly. Re-run Cells 5 and 6 to reinitialize, then re-run Cells 7-9.
  • Help me debug my training loop

✔️ Awesome, I've got everything!

Great. Make sure your notebook is saved before moving on.

ⓧ I'd like to double check the full code

Here is everything your notebook should contain up to this point:

# Cell 1: Imports
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt

# Cell 2: Load Fashion-MNIST dataset
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)

class_names = [
    "T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",
    "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"
]

print(f"Training set size: {len(training_data)}")
print(f"Test set size: {len(test_data)}")
print(f"Image shape: {training_data[0][0].shape}")
print(f"Number of classes: {len(class_names)}")

# Cell 3: Visualize sample images
fig, axes = plt.subplots(2, 5, figsize=(12, 5))
for i, ax in enumerate(axes.flat):
    image, label = training_data[i]
    ax.imshow(image.squeeze(), cmap="gray")
    ax.set_title(class_names[label])
    ax.axis("off")
plt.tight_layout()
plt.show()

# Cell 4: Create DataLoaders
batch_size = 64

train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=False)

# Cell 5: Define the neural network
class FashionClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.network = nn.Sequential(
            nn.Linear(28 * 28, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 10),
        )

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

model = FashionClassifier()
print(model)

# Cell 6: Set up loss function and optimizer
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# Cell 7: Define training function
def train_one_epoch(dataloader, model, loss_fn, optimizer):
    model.train()
    running_loss = 0.0
    for batch, (X, y) in enumerate(dataloader):
        pred = model(X)
        loss = loss_fn(pred, y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    avg_loss = running_loss / len(dataloader)
    return avg_loss

# Cell 8: Define evaluation function
def evaluate(dataloader, model, loss_fn):
    model.eval()
    correct = 0
    total = 0
    test_loss = 0.0

    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).sum().item()
            total += y.size(0)

    accuracy = correct / total
    avg_loss = test_loss / len(dataloader)
    return accuracy, avg_loss

# Cell 9: Train the model
epochs = 10

for epoch in range(epochs):
    train_loss = train_one_epoch(train_dataloader, model, loss_fn, optimizer)
    accuracy, test_loss = evaluate(test_dataloader, model, loss_fn)
    print(f"Epoch {epoch+1:>2d} | Train Loss: {train_loss:.4f} | Test Loss: {test_loss:.4f} | Accuracy: {accuracy*100:.1f}%")

print("\nTraining complete!")

Your model just went from random guessing to correctly identifying most clothing items. Next up, you will visualize what the model gets right and wrong to understand its strengths and weaknesses.

Visualize Predictions

Your model just finished 10 epochs of training and can classify clothing images at roughly 85% accuracy. That single number is encouraging, but it hides important details about where the model succeeds and where it struggles.

In this step, you'll visualize individual predictions on real test images and break down accuracy by clothing category. This reveals whether the model confuses visually similar garments like shirts and coats.

In this step, get ready to:
  • Display a grid of predictions with color-coded correct and incorrect labels.
  • Analyze which clothing categories confuse the model.
  • Print per-class accuracy to quantify strengths and weaknesses.
Write the prediction visualization function

A single accuracy percentage tells you the model works, but not how it fails. Showing the model's guesses on actual images makes its mistakes visible and interpretable.

  • Create a new cell at the bottom of your notebook (Cell 10) and add the following prediction visualization code:
# Cell 10: Visualize predictions
def visualize_predictions(model, test_data, class_names, num_images=15):
    model.eval()
    fig, axes = plt.subplots(3, 5, figsize=(14, 8))

    for i, ax in enumerate(axes.flat):
        image, true_label = test_data[i]
        with torch.no_grad():
            pred = model(image.unsqueeze(0))
            pred_label = pred.argmax(1).item()

        ax.imshow(image.squeeze(), cmap="gray")
        color = "green" if pred_label == true_label else "red"
        ax.set_title(
            f"Pred: {class_names[pred_label]}\nTrue: {class_names[true_label]}",
            color=color,
            fontsize=9,
        )
        ax.axis("off")

    plt.tight_layout()
    plt.show()

visualize_predictions(model, test_data, class_names)

What does this code do?

  • model.eval() switches the model to evaluation mode so it makes predictions without updating weights.
  • image.unsqueeze(0) adds a batch dimension because the model expects batches of images, not a single image.
  • pred.argmax(1).item() picks the class with the highest prediction score and converts it to a plain Python integer.
  • The color variable sets the title green for correct predictions and red for incorrect ones, so mistakes stand out immediately.
  • Run the cell by pressing Shift+Enter.
  • You should see a 3x5 grid of grayscale clothing images. Each image has two lines of text above it showing the predicted label and the true label. Correct predictions appear in green and incorrect ones appear in red.

Not seeing the image grid?

  • Make sure you ran all previous cells first (Cells 1 through 9). The model needs to be trained before it can make predictions.
  • Check that you have import matplotlib.pyplot as plt in Cell 1. Without matplotlib, no plots render.
  • help me debug why my prediction visualization isn't showing up
Analyze which clothing categories get confused

Look at the grid you just generated. Most titles should be green, confirming the model gets the majority right.

Focus on the red titles. These are the model's mistakes. Notice which categories appear in the incorrect predictions.

Why does the model confuse certain items?

Your model sees each image as a flat list of 784 pixel values. It has no concept of shape, texture, or structure. Items that share similar silhouettes (like Shirts and T-shirts, or Coats and Pullovers) produce similar pixel patterns.

This is a fundamental limitation of fully-connected networks. Convolutional neural networks (CNNs) solve this by learning spatial features, but that's a more advanced architecture for another project.

Print per-class accuracy to identify strengths and weaknesses

The visualization grid shows a small sample. To get a complete picture, you need to check accuracy across all 10,000 test images broken down by category.

  • Create a new cell (Cell 11) and add the per-class accuracy function:
# Cell 11: Per-class accuracy
def per_class_accuracy(dataloader, model, class_names):
    model.eval()
    class_correct = [0] * 10
    class_total = [0] * 10

    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            predicted = pred.argmax(1)
            for i in range(len(y)):
                label = y[i].item()
                class_total[label] += 1
                if predicted[i] == label:
                    class_correct[label] += 1

    print("\nPer-class accuracy:")
    print("-" * 35)
    for i in range(10):
        acc = class_correct[i] / class_total[i] * 100
        print(f"{class_names[i]:<12s}: {acc:.1f}%")

per_class_accuracy(test_dataloader, model, class_names)

What does this code do?

  • class_correct and class_total are lists of 10 counters (one per clothing category) that track how many the model got right versus how many it saw.
  • The function loops through every batch in the test set, compares each prediction against the true label, and increments the correct counter when they match.
  • At the end, it divides correct by total for each class and prints the percentage. This tells you exactly which categories the model handles well and which it struggles with.
  • Run the cell by pressing Shift+Enter.
  • You should see a table listing all 10 clothing categories with their individual accuracy percentages. Categories like Trouser and Sneaker should score above 90%, while Shirt is likely the lowest.

Seeing a division error or all zeros?

  • Make sure you passed test_dataloader (not test_data) as the first argument. The function expects a DataLoader that yields batches.
  • Verify Cell 4 was run so that test_dataloader exists in your notebook's memory.
  • help me debug my per-class accuracy function

✔️ Awesome, I've got everything!

Great. Double check you've saved your notebook file.

ⓧ I'd like to double check the full code

# Cell 1: Imports
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt

# Cell 2: Load Fashion-MNIST dataset
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)

class_names = [
    "T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",
    "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"
]

print(f"Training set size: {len(training_data)}")
print(f"Test set size: {len(test_data)}")
print(f"Image shape: {training_data[0][0].shape}")
print(f"Number of classes: {len(class_names)}")

# Cell 3: Visualize sample images
fig, axes = plt.subplots(2, 5, figsize=(12, 5))
for i, ax in enumerate(axes.flat):
    image, label = training_data[i]
    ax.imshow(image.squeeze(), cmap="gray")
    ax.set_title(class_names[label])
    ax.axis("off")
plt.tight_layout()
plt.show()

# Cell 4: Create DataLoaders
batch_size = 64

train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=False)

# Cell 5: Define the neural network
class FashionClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.network = nn.Sequential(
            nn.Linear(28 * 28, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 10),
        )

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

model = FashionClassifier()
print(model)

# Cell 6: Set up loss function and optimizer
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# Cell 7: Define training function
def train_one_epoch(dataloader, model, loss_fn, optimizer):
    model.train()
    running_loss = 0.0
    for batch, (X, y) in enumerate(dataloader):
        pred = model(X)
        loss = loss_fn(pred, y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    avg_loss = running_loss / len(dataloader)
    return avg_loss

# Cell 8: Define evaluation function
def evaluate(dataloader, model, loss_fn):
    model.eval()
    correct = 0
    total = 0
    test_loss = 0.0

    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).sum().item()
            total += y.size(0)

    accuracy = correct / total
    avg_loss = test_loss / len(dataloader)
    return accuracy, avg_loss

# Cell 9: Train the model
epochs = 10

for epoch in range(epochs):
    train_loss = train_one_epoch(train_dataloader, model, loss_fn, optimizer)
    accuracy, test_loss = evaluate(test_dataloader, model, loss_fn)
    print(f"Epoch {epoch+1:>2d} | Train Loss: {train_loss:.4f} | Test Loss: {test_loss:.4f} | Accuracy: {accuracy*100:.1f}%")

print("\nTraining complete!")

# Cell 10: Visualize predictions
def visualize_predictions(model, test_data, class_names, num_images=15):
    model.eval()
    fig, axes = plt.subplots(3, 5, figsize=(14, 8))

    for i, ax in enumerate(axes.flat):
        image, true_label = test_data[i]
        with torch.no_grad():
            pred = model(image.unsqueeze(0))
            pred_label = pred.argmax(1).item()

        ax.imshow(image.squeeze(), cmap="gray")
        color = "green" if pred_label == true_label else "red"
        ax.set_title(
            f"Pred: {class_names[pred_label]}\nTrue: {class_names[true_label]}",
            color=color,
            fontsize=9,
        )
        ax.axis("off")

    plt.tight_layout()
    plt.show()

visualize_predictions(model, test_data, class_names)

# Cell 11: Per-class accuracy
def per_class_accuracy(dataloader, model, class_names):
    model.eval()
    class_correct = [0] * 10
    class_total = [0] * 10

    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            predicted = pred.argmax(1)
            for i in range(len(y)):
                label = y[i].item()
                class_total[label] += 1
                if predicted[i] == label:
                    class_correct[label] += 1

    print("\nPer-class accuracy:")
    print("-" * 35)
    for i in range(10):
        acc = class_correct[i] / class_total[i] * 100
        print(f"{class_names[i]:<12s}: {acc:.1f}%")

per_class_accuracy(test_dataloader, model, class_names)

Secret mission

Deep Dive into Model Performance

Your model hits ~85% accuracy, but which clothing categories is it confusing? Build a confusion matrix to find out, then design a deeper network with dropout regularization to push accuracy higher.

Clean Up Your Resources

Clean Up Your Resources

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

Resources you used:

  • Project folder on your Desktop containing fashion_mnist_classifier.ipynb and the data/ directory (Fashion-MNIST dataset).
  • Python 3.12 installed via pyenv.
  • pyenv (Python version manager) installed at ~/.pyenv.
  • Jupyter extension installed in VS Code.

Keep everything running

No action needed. Choose this if you plan to keep experimenting with your neural network or want to revisit the notebook later.

  • Your project folder stays on your Desktop, ready for further experimentation.
  • pyenv and Python 3.12 remain installed for future Python projects.
  • The Jupyter extension stays active in VS Code for other notebooks.

Pause - I'll come back to this later

Shut down any running Jupyter kernels to free up memory. Your files and tools stay installed for when you return.

  • In VS Code, open the notebook and click Restart in the kernel toolbar to shut down the running kernel.
  • Close the notebook tab in VS Code.
  • If you activated a virtual environment in your terminal, deactivate it by running:
deactivate

Everything else stays in place. When you come back, reactivate the virtual environment and reopen the notebook.

Delete - I don't want to use this again

Remove all project resources and tools installed during this project.

  • Delete the project folder from your Desktop by running:
cd ~/Desktop
rm -rf fashion-mnist-classifier
  • Remove pyenv and the Python 3.12 installation by running:
rm -rf ~/.pyenv
  • Open your shell configuration file (~/.bashrc or ~/.zshrc) and remove the three lines that pyenv added during installation. They look like this:
export PYENV_ROOT="$HOME/.pyenv"
[[ -d $PYENV_ROOT/bin ]] && export PATH="$PYENV_ROOT/bin:$PATH"
eval "$(pyenv init -)"
  • Uninstall the Jupyter extension in VS Code. Open the Extensions sidebar (Ctrl+Shift+X), find Jupyter, and click Uninstall.

Nice Work!

Nice Work!

You just trained a neural network from scratch. Starting from raw pixel data and random weights, you built a model that learned to recognize clothing items with ~85% accuracy.

You've learned how to:

  • Built a complete ML pipeline in PyTorch covering data loading, model definition, training, and evaluation.
  • Trained a neural network that classifies 10 types of clothing at ~85% accuracy using backpropagation and SGD with momentum.
  • Visualized model predictions and identified failure patterns using per-class accuracy analysis.
  • Built a deeper network with dropout regularization and compared architectures using confusion matrices.

Ready to quiz yourself?