Build a CNN Image Classifier with PyTorch

Build and train a convolutional neural network to classify images with PyTorch.

Introduction

30 Second Summary

Your phone can tell the difference between a cat and a dog in photos. Behind that ability is a type of AI called a convolutional neural network, trained on thousands of labeled images.

In this project, you will build your own image classifier from scratch using PyTorch. You will construct a CNN, train it on 60,000 real-world color images, and evaluate how well it can recognize objects like airplanes, cars, and birds.

What You'll Build

A Python-based convolutional neural network that classifies 32x32 color images into 10 categories using PyTorch and the CIFAR-10 dataset.

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

  • A complete CNN with three convolutional blocks and two fully connected layers, trained from scratch in a single Python script.
  • Training loss charts and per-class accuracy breakdowns showing how your model performs across all 10 image categories.
  • Visual prediction results with color-coded labels showing where your model gets it right and where it struggles.
  • Secret Mission: Supercharge your CNN with Batch Normalization, Dropout, and data augmentation to boost accuracy.

Are there any prerequisites?

You should have basic Python knowledge (variables, functions, loops) and some familiarity with machine learning concepts like training and testing. No prior deep learning or PyTorch experience is needed.

Set Up PyTorch

Every deep learning project starts with the right foundation. PyTorch is the open-source framework you will use to build, train, and evaluate your image classifier.

In this step, you will create a dedicated project folder, set up a Python virtual environment to keep your packages isolated, and install PyTorch alongside matplotlib for visualization.

In this step, get ready to:
  • Create a project folder and a Python virtual environment.
  • Install PyTorch (CPU) and matplotlib.
  • Verify everything works by running a test script.
Create a project folder and virtual environment

A virtual environment keeps this project's Python packages completely separate from any other Python projects on your machine. You will create one inside a dedicated project folder on your Desktop.

  • Press Cmd+` (macOS) or Ctrl+` (Windows) to open the terminal panel in VS Code.
  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder by running this command:
mkdir cifar10-classifier
  • Now open this folder in VS Code.
  • Click File in the top menu bar.
  • Click Open Folder.
  • Navigate to your Desktop, select the cifar10-classifier folder, then click Open.

Trust this folder?

If VS Code asks whether you trust the authors, click Yes, I trust the authors. This is your own project folder, so it is safe to trust.

  • Press Cmd+` (macOS) or Ctrl+` (Windows) to open a new terminal panel.

The terminal opens inside cifar10-classifier automatically.

  • Create a virtual environment and activate it by running these commands:

macOS

python3 -m venv venv
source venv/bin/activate

Windows

python -m venv venv
venv\Scripts\activate

What does this do?

The first command creates an isolated Python environment inside a folder called venv. The second command activates that environment so every package you install stays contained in this project.

Your terminal prompt now starts with (venv), confirming the virtual environment is active.

Don't see (venv) in your prompt?

  • Make sure you ran the activation command that matches your operating system. macOS uses source while Windows uses the Scripts folder path.
  • On Windows, you may see an error about script execution policies. Help me fix a script execution policy error in PowerShell.
Install PyTorch and matplotlib

PyTorch handles tensor operations, neural network layers, and automatic differentiation. matplotlib creates the charts you will use to visualize training progress later.

  • Install PyTorch and torchvision by running this command:
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu

What does --index-url do?

The --index-url flag tells pip to download PyTorch from a special CPU-only repository instead of the default PyPI index. This keeps the download smaller and avoids GPU dependencies you do not need for this project.

The install takes a minute or two. You should see a "Successfully installed" message listing torch and torchvision.

Install command failing?

  • Make sure (venv) appears in your terminal prompt. If it disappeared, reactivate your virtual environment using the command from the previous substep.
  • Check your internet connection. PyTorch is a large download.
  • Help me troubleshoot a PyTorch installation error.
  • Install matplotlib by running:
pip install matplotlib

Why a separate install?

The --index-url flag in the previous command replaced the default package source entirely. matplotlib comes from the standard PyPI index, so it needs its own pip install command.

You should see a "Successfully installed" message for matplotlib and its dependencies.

Verify the installation

Time to confirm everything is wired up correctly with a quick test script.

  • In the VS Code Explorer sidebar on the left, click the New File icon (the page icon with a plus sign) next to CIFAR10-CLASSIFIER.
  • Type classifier.py and press Enter.
  • Add the following code to classifier.py:
import torch
print(f'PyTorch version: {torch.__version__}')
x = torch.rand(3, 3)
print(x)

What does this code do?

  • import torch loads the PyTorch library into your script.
  • torch.__version__ prints which version of PyTorch you installed, confirming the package is available.
  • torch.rand(3, 3) creates a 3x3 grid of random numbers called a tensor. If this prints without errors, PyTorch's core operations are working.
  • Save classifier.py.
  • Run the script by executing this command in your terminal:
python classifier.py

You should see two things: a PyTorch version starting with 2.7 and a 3x3 grid of random decimal numbers.

Seeing an error?

  • If you see an error about the torch module not being found, your virtual environment may not be active. Check for (venv) in your terminal prompt and reactivate if needed.
  • If you see a "file not found" error, make sure you saved classifier.py inside the cifar10-classifier folder.
  • Help me debug an error when running classifier.py.

✔️ Awesome, I've got everything!

Great. Double check you have saved classifier.py before moving on.

ⓧ I'd like to double check the full code

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

import torch
print(f'PyTorch version: {torch.__version__}')
x = torch.rand(3, 3)
print(x)

Your environment is ready and PyTorch is installed. Next up, you will load the CIFAR-10 dataset and explore the images your neural network will learn to classify.

Load and Explore the CIFAR-10 Dataset

Your PyTorch environment is set up and verified. Now your model needs data to learn from.

CIFAR-10 is a benchmark dataset of 60,000 color images across 10 categories like airplanes, cars, and cats. In this step, you'll load it into your project and visualize some samples to see exactly what your CNN will be classifying.

In this step, get ready to:
  • Load the CIFAR-10 training and test sets using torchvision.
  • Visualize a batch of sample images to see what the dataset looks like.
Load the dataset with torchvision

Every image in CIFAR-10 is a tiny 32x32 pixel color photograph. Before your CNN can process these images, they need to be converted into tensors and normalized so pixel values fall within a consistent range.

Why CIFAR-10 instead of MNIST?

MNIST (handwritten digits) is often the "hello world" dataset for machine learning, but it's grayscale and simple enough that a basic fully connected network can solve it. CIFAR-10 uses 3-channel color images of real-world objects, which actually require convolutional processing to classify well.

This makes CIFAR-10 a much better fit for learning how CNNs work. The complexity justifies the architecture.

  • In classifier.py, select all the existing code by pressing Cmd+A (macOS) or Ctrl+A (Windows) and delete it.
  • Set up the imports and data preprocessing transform by pasting this code:
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt


# --- Data Loading ---
# Chain preprocessing steps: convert images to tensors, then normalize
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

What does this code do?

The imports bring in everything you'll need across this project. torch is the core library. torchvision provides datasets and image transforms. matplotlib handles plotting.

transforms.Compose chains two preprocessing steps together. ToTensor() converts each image from a PIL image into a PyTorch tensor. Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) scales pixel values from the [0, 1] range down to [-1, 1], which helps the model train more effectively.

You won't see a change yet. The next chunk loads the actual dataset using this transform.

  • Add the dataset loading code below the transform definition by pasting this code at the bottom of the file:
# Download and load the CIFAR-10 training set
trainset = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=transform
)
trainloader = DataLoader(trainset, batch_size=64, shuffle=True, num_workers=0)

# Download and load the CIFAR-10 test set
testset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=transform
)
testloader = DataLoader(testset, batch_size=64, shuffle=False, num_workers=0)

# Human-readable names for the 10 CIFAR-10 categories
classes = ('airplane', 'automobile', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

print(f'{len(trainset)} training images, {len(testset)} test images')

What does this code do?

  • torchvision.datasets.CIFAR10 downloads the dataset the first time you run the script. The download=True flag tells torchvision to fetch it automatically if it's not already on disk.
  • DataLoader wraps each dataset so you can iterate through it in batches of 64 images at a time, instead of one image at a time.
  • num_workers=0 keeps data loading on the main process. This avoids multiprocessing errors that can occur on Windows.
  • classes maps numeric labels (0 through 9) to human-readable category names like airplane and truck.
  • Save classifier.py.
  • Run the script by entering this command in your terminal:
python classifier.py
  • The first run downloads CIFAR-10 (~163 MB), so it may take a minute depending on your internet speed. You should see download progress in the terminal, followed by the line 50000 training images, 10000 test images.

Download failing or hanging?

  • Check your internet connection. The dataset is downloaded from an external server and requires a stable connection.
  • If you see ModuleNotFoundError: No module named 'torchvision', make sure your virtual environment is activated. You should see (venv) at the start of your terminal prompt.
  • If the download stalls, press Ctrl+C to cancel and run the script again. Torchvision picks up where it left off.

Why is my CIFAR-10 download failing or showing a ModuleNotFoundError?

Visualize sample images

Numbers on a screen don't tell you much about what your data actually looks like. Plotting a few sample images helps you confirm the data loaded correctly and gives you a feel for what your CNN will be working with.

  • Add the visualization function at the bottom of classifier.py by pasting this code below the print statement:

# --- Visualize Samples ---
def show_images(loader):
    # Grab one batch of images from the loader
    dataiter = iter(loader)
    images, labels = next(dataiter)

    # Display 8 images in a row
    fig, axes = plt.subplots(1, 8, figsize=(12, 2))
    for i in range(8):
        # Reverse the normalization (multiply by 0.5 and add 0.5)
        img = images[i] * 0.5 + 0.5
        # Rearrange from CHW (channels, height, width) to HWC for matplotlib
        axes[i].imshow(img.permute(1, 2, 0).numpy())
        axes[i].set_title(classes[labels[i]], fontsize=8)
        axes[i].axis('off')
    plt.tight_layout()
    plt.savefig('sample_images.png', dpi=150)
    plt.show()


show_images(trainloader)

What does this code do?

  • iter(loader) and next(dataiter) grab the first batch of 64 images from the training loader.
  • images[i] * 0.5 + 0.5 reverses the normalization you applied earlier, bringing pixel values back to the [0, 1] range so matplotlib can display them correctly.
  • img.permute(1, 2, 0) rearranges the image dimensions. PyTorch stores images as (channels, height, width), but matplotlib expects (height, width, channels).
  • plt.savefig('sample_images.png', dpi=150) saves the figure to disk so you have a permanent record of what the data looks like.
  • Save classifier.py and run the script again:
python classifier.py
  • A matplotlib window pops up showing 8 labeled CIFAR-10 images. You'll see tiny color photographs of objects like frogs, ships, or trucks with their category names above each one.
  • Close the matplotlib window to let the script finish. Confirm that sample_images.png was saved to your cifar10-classifier folder by running:
ls
  • You should see sample_images.png, classifier.py, and the data folder (where CIFAR-10 was downloaded) in the output.

Not seeing the images or sample_images.png?

  • If the matplotlib window doesn't appear, don't worry. The plt.savefig() call still saves the image to disk. Open sample_images.png from your file explorer to verify it.
  • If you see a NameError: name 'classes' is not defined, make sure the classes tuple is defined above the show_images function.
  • Check that you're calling show_images(trainloader) at the bottom of the file, outside the function definition.

Why aren't my CIFAR-10 sample images displaying or saving correctly?

✔️ Awesome, I've got everything!

Great. Double check you've saved classifier.py before moving on.

ⓧ I'd like to double check the full code

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt


# --- Data Loading ---
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

trainset = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=transform
)
trainloader = DataLoader(trainset, batch_size=64, shuffle=True, num_workers=0)

testset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=transform
)
testloader = DataLoader(testset, batch_size=64, shuffle=False, num_workers=0)

classes = ('airplane', 'automobile', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

print(f'{len(trainset)} training images, {len(testset)} test images')


# --- Visualize Samples ---
def show_images(loader):
    dataiter = iter(loader)
    images, labels = next(dataiter)

    fig, axes = plt.subplots(1, 8, figsize=(12, 2))
    for i in range(8):
        img = images[i] * 0.5 + 0.5
        axes[i].imshow(img.permute(1, 2, 0).numpy())
        axes[i].set_title(classes[labels[i]], fontsize=8)
        axes[i].axis('off')
    plt.tight_layout()
    plt.savefig('sample_images.png', dpi=150)
    plt.show()


show_images(trainloader)

Your data is loaded and you've seen what CIFAR-10 looks like up close. Next up, you'll build the CNN architecture that will learn to classify these images.

Build the CNN Architecture

You've loaded 50,000 training images and seen what CIFAR-10 looks like up close. Now it's time to build the neural network that will learn to classify those images.

A convolutional neural network processes images through layers that detect increasingly complex patterns. Early layers spot edges and color gradients. Deeper layers combine those into recognizable features like wheels, wings, or fur.

In this step, get ready to:
  • Define a CNN class with three convolutional blocks and two fully connected layers.
  • Instantiate the model and inspect its architecture and parameter count.
Define the CNN class

In PyTorch, neural networks are defined as classes that extend nn.Module. Your class needs two methods: __init__ to set up the layers, and forward to define how data flows through them.

  • Add the ImageClassifier class to the bottom of classifier.py, below the show_images(trainloader) line:
# --- CNN Architecture ---
class ImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(128 * 4 * 4, 256)
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

What does this code do?

__init__ defines three Conv2d layers that progressively extract features from the image. The channels increase from 3 (one per RGB color) to 32, then 64, then 128 feature maps. All three use 3x3 kernels with padding=1 to preserve spatial dimensions before pooling.

MaxPool2d(2, 2) halves the width and height after each convolution. The 32x32 input shrinks to 16x16, then 8x8, then 4x4.

After flattening the 128 feature maps at 4x4 into a single vector of 2048 values, fc1 compresses them to 256 features. fc2 maps those to 10 output scores, one per CIFAR-10 class.

forward defines the data path: convolution, ReLU activation, then pooling, repeated three times. After flattening, data passes through the two fully connected layers.

  • Save classifier.py.
  • You won't see a change yet. The next substep instantiates the model and prints its architecture.

Seeing a red underline in your editor?

Make sure every line inside the class is indented with exactly four spaces. Python uses indentation to define class and method bodies.

Check that __init__ has two underscores on each side, not one.

Why is my ImageClassifier class showing a syntax error?

Instantiate the model and inspect it

The class defines the blueprint. Now you need to create an actual model instance from it and place it on the right computing device.

  • Add the following code below the ImageClassifier class, at the bottom of classifier.py:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')

model = ImageClassifier().to(device)
print(model)

total_params = sum(p.numel() for p in model.parameters())
print(f'Total parameters: {total_params:,}')

What does this code do?

torch.device checks whether a CUDA-capable GPU is available. Since you installed CPU-only PyTorch, this defaults to cpu.

ImageClassifier().to(device) creates an instance of your CNN and moves all its parameters to the selected device.

p.numel() counts the elements in each parameter tensor. Summing across all parameters gives the total number of learnable values in the network.

  • Save classifier.py.
  • Run the script to see your model's architecture by running this command:
python classifier.py
  • You should see Using device: cpu, the full model architecture listing all six layers, and Total parameters: showing a count of approximately 620,000.

Seeing a NameError or unexpected output?

If you see NameError: name 'ImageClassifier' is not defined, make sure the instantiation code sits below the class definition, not above it.

If you see an IndentationError, check that the class body uses consistent four-space indentation throughout. Mixing tabs and spaces causes this error.

Why am I getting an error when running my CNN script?

✔️ Awesome, I've got everything!

Great. Double check you've saved classifier.py.

ⓧ I'd like to double check the full code

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt


# --- Data Loading ---
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

trainset = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=transform
)
trainloader = DataLoader(trainset, batch_size=64, shuffle=True, num_workers=0)

testset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=transform
)
testloader = DataLoader(testset, batch_size=64, shuffle=False, num_workers=0)

classes = ('airplane', 'automobile', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

print(f'{len(trainset)} training images, {len(testset)} test images')


# --- Visualize Samples ---
def show_images(loader):
    dataiter = iter(loader)
    images, labels = next(dataiter)

    fig, axes = plt.subplots(1, 8, figsize=(12, 2))
    for i in range(8):
        img = images[i] * 0.5 + 0.5
        axes[i].imshow(img.permute(1, 2, 0).numpy())
        axes[i].set_title(classes[labels[i]], fontsize=8)
        axes[i].axis('off')
    plt.tight_layout()
    plt.savefig('sample_images.png', dpi=150)
    plt.show()


show_images(trainloader)


# --- CNN Architecture ---
class ImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(128 * 4 * 4, 256)
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')

model = ImageClassifier().to(device)
print(model)

total_params = sum(p.numel() for p in model.parameters())
print(f'Total parameters: {total_params:,}')

Your CNN architecture is locked in and ready to learn. Next up, you'll train this model on all 50,000 images and watch its loss decrease over time.

Train the Model

Your CNN architecture is built and ready. Now it's time to teach it how to recognize images.

Training a neural network means showing it thousands of labeled images and adjusting its internal weights each time it makes a wrong prediction. In this step, you'll set up the tools that guide this learning process and write the loop that runs it.

In this step, get ready to:
  • Set up a loss function and optimizer to guide how the model learns.
  • Write a training loop that teaches the model over 10 epochs.
  • Plot the training loss to visualize the model's progress.
Set up the loss function and optimizer

Before the model can learn, it needs two things. A loss function measures how wrong the model's predictions are. An optimizer decides how to adjust the model's weights to reduce that error.

  • Add the following code to the bottom of classifier.py, below the line that prints Total parameters:
# --- Training ---
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

num_epochs = 10
train_losses = []

What does this code do?

  • nn.CrossEntropyLoss() is the standard loss function for multi-class classification. It calculates how far the model's predicted probabilities are from the correct answer.
  • optim.Adam is an optimizer that adapts the learning rate for each parameter individually. It typically converges faster than basic SGD.
  • lr=0.001 sets the learning rate, which controls how big each weight update is. Too high and the model overshoots. Too low and it learns too slowly.
  • train_losses is an empty list that will store the loss value from each epoch so you can plot it later.

You won't see any changes yet. The next chunk uses these to actually train the model.

Write the training loop

The training loop is where the model actually learns. Each pass through the entire dataset is called an epoch. Within each epoch, the model processes images in small batches, measures its error, and adjusts its weights to do better next time.

  • Add the training loop below the train_losses = [] line:
for epoch in range(num_epochs):
    running_loss = 0.0
    model.train()

    for inputs, labels in trainloader:
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    epoch_loss = running_loss / len(trainloader)
    train_losses.append(epoch_loss)
    print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {epoch_loss:.4f}')

print('Training complete!')

What's happening inside the loop?

Each epoch repeats four core steps for every batch of 64 images:

  • optimizer.zero_grad() resets the gradients from the previous batch. Without this, gradients accumulate and corrupt the updates.
  • outputs = model(inputs) is the forward pass. Images flow through all the conv and linear layers to produce predictions.
  • loss.backward() is backpropagation. It calculates how much each weight contributed to the error.
  • optimizer.step() updates every weight in the model based on those gradients.

You won't see output yet. Add the plotting code next, then you'll run everything together.

Plot the training loss

A loss chart lets you visually confirm that the model is actually learning. If the loss decreases over epochs, the model is getting better at its predictions.

  • Add the plotting code below the print('Training complete!') line:
# --- Plot Training Loss ---
plt.figure(figsize=(8, 4))
plt.plot(range(1, num_epochs + 1), train_losses, marker='o')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training Loss Over Epochs')
plt.grid(True)
plt.savefig('training_loss.png', dpi=150)
plt.show()

What does this code do?

  • plt.plot() draws a line chart with epochs on the x-axis and loss values on the y-axis.
  • plt.savefig('training_loss.png') saves the chart as an image file in your project folder so you can review it later.
  • Save classifier.py.
  • Run the script by entering this command. Training takes approximately 5-10 minutes on CPU:
python classifier.py

What should I see?

First, the script reloads the dataset and displays sample images (same as before). Then training begins.

You'll see loss values printing for each epoch. The loss should start around 1.4-1.5 in epoch 1 and decrease to around 0.6-0.8 by epoch 10. After training completes, a chart window appears and training_loss.png is saved to your cifar10-classifier folder.

Training seems stuck or loss isn't decreasing?

  • Each epoch processes 782 batches of 64 images. On CPU, a single epoch can take 30-60 seconds. Wait for the first Epoch 1/10 line to print before concluding something is wrong.
  • If you see a NameError for criterion or optimizer, make sure the training setup code (the # --- Training --- section) appears before the training loop.
  • If the matplotlib chart window blocks the script, close the chart window to let the script finish.

Help me debug my training loop

✔️ Awesome, I've got everything!

Great. Double check you've saved your file before moving on.

ⓧ I'd like to double check the full code

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt


# --- Data Loading ---
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

trainset = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=transform
)
trainloader = DataLoader(trainset, batch_size=64, shuffle=True, num_workers=0)

testset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=transform
)
testloader = DataLoader(testset, batch_size=64, shuffle=False, num_workers=0)

classes = ('airplane', 'automobile', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

print(f'{len(trainset)} training images, {len(testset)} test images')


# --- Visualize Samples ---
def show_images(loader):
    dataiter = iter(loader)
    images, labels = next(dataiter)

    fig, axes = plt.subplots(1, 8, figsize=(12, 2))
    for i in range(8):
        img = images[i] * 0.5 + 0.5
        axes[i].imshow(img.permute(1, 2, 0).numpy())
        axes[i].set_title(classes[labels[i]], fontsize=8)
        axes[i].axis('off')
    plt.tight_layout()
    plt.savefig('sample_images.png', dpi=150)
    plt.show()


show_images(trainloader)


# --- CNN Architecture ---
class ImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(128 * 4 * 4, 256)
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')

model = ImageClassifier().to(device)
print(model)

total_params = sum(p.numel() for p in model.parameters())
print(f'Total parameters: {total_params:,}')


# --- Training ---
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

num_epochs = 10
train_losses = []

for epoch in range(num_epochs):
    running_loss = 0.0
    model.train()

    for inputs, labels in trainloader:
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    epoch_loss = running_loss / len(trainloader)
    train_losses.append(epoch_loss)
    print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {epoch_loss:.4f}')

print('Training complete!')


# --- Plot Training Loss ---
plt.figure(figsize=(8, 4))
plt.plot(range(1, num_epochs + 1), train_losses, marker='o')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training Loss Over Epochs')
plt.grid(True)
plt.savefig('training_loss.png', dpi=150)
plt.show()

Your model has learned to recognize patterns in 50,000 images. Next up, you'll put it to the test and find out exactly how accurate it is across all 10 categories.

Evaluate and Visualize Predictions

Your model is trained and the loss curve shows it learned something. But a decreasing loss only tells you the model is getting better at the training data. The real question is: how well does it handle images it has never seen before?

In this step, you will evaluate your CNN on the 10,000 test images it was never trained on. You will measure overall accuracy, break it down by class, and visualize the model's actual predictions with color-coded labels.

In this step, get ready to:
  • Compute overall test accuracy across all 10,000 test images.
  • Compute per-class accuracy to see which categories your model handles best.
  • Visualize predictions on test images and save the trained model to disk.
Compute overall test accuracy

During training, your model saw every image in the training set multiple times. To measure real performance, you need to test it on images it has never encountered. This is what the test set is for.

  • In classifier.py, scroll to the bottom of the file (after the training loss plot code).
  • Add the evaluation code at the end of the file:

# --- Evaluation ---
model.eval()
correct = 0
total = 0

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

print(f'Test Accuracy: {100 * correct / total:.2f}%')

What does this code do?

model.eval() switches the model from training mode to evaluation mode. This disables behaviors like dropout that only apply during training.

torch.no_grad() tells PyTorch to skip computing gradients. Since you are only making predictions (not training), this saves memory and speeds things up.

torch.max(outputs, 1) returns the class with the highest confidence score for each image. The _ discards the score values since you only need the predicted class indices.

You won't see a change yet. The next chunk adds per-class accuracy before you run the script.

Compute per-class accuracy

Overall accuracy is a useful single number, but it hides interesting details. Some categories (like trucks and ships) are easier for a CNN to recognize than others (like cats and dogs). A per-class breakdown reveals where your model excels and where it struggles.

  • Add the per-class accuracy code below the overall accuracy print statement:

# --- Per-Class Accuracy ---
class_correct = {c: 0 for c in classes}
class_total = {c: 0 for c in classes}

with torch.no_grad():
    for inputs, labels in testloader:
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)
        _, predicted = torch.max(outputs, 1)
        for i in range(labels.size(0)):
            label = classes[labels[i]]
            class_total[label] += 1
            if predicted[i] == labels[i]:
                class_correct[label] += 1

print('\nPer-Class Accuracy:')
for cls in classes:
    accuracy = 100 * class_correct[cls] / class_total[cls]
    print(f'  {cls:12s}: {accuracy:.1f}%')

What does this code do?

Two dictionaries track the results. class_correct counts how many images of each class the model got right. class_total counts how many images of each class exist in the test set.

The loop iterates over every test image, checks whether the prediction matches the true label, and tallies the results per class. Dividing correct by total gives each class its own accuracy percentage.

You won't see a change yet. The next chunk adds prediction visualization and saves the model, then you will run the full script.

Visualize predictions and save the model

Numbers are useful, but seeing your model's predictions on actual images is far more revealing. A visualization lets you spot patterns in what the model gets right and wrong.

  • Add the prediction visualization function and model save code at the end of classifier.py:

# --- Visualize Predictions ---
def show_predictions(model, loader):
    dataiter = iter(loader)
    images, labels = next(dataiter)

    model.eval()
    with torch.no_grad():
        outputs = model(images.to(device))
        _, predicted = torch.max(outputs, 1)

    fig, axes = plt.subplots(2, 8, figsize=(14, 4))
    for i in range(16):
        row = i // 8
        col = i % 8
        img = images[i] * 0.5 + 0.5
        axes[row, col].imshow(img.permute(1, 2, 0).numpy())
        color = 'green' if predicted[i] == labels[i] else 'red'
        axes[row, col].set_title(classes[predicted[i]], fontsize=8, color=color)
        axes[row, col].axis('off')
    plt.suptitle('Green = Correct, Red = Wrong')
    plt.tight_layout()
    plt.savefig('predictions.png', dpi=150)
    plt.show()


show_predictions(model, testloader)


# --- Save Model ---
torch.save(model.state_dict(), 'cifar10_model.pth')
print('Model saved to cifar10_model.pth')

What does this code do?

show_predictions grabs a batch of 64 test images and runs them through the model. It displays the first 16 in a 2x8 grid. Each title shows the predicted class name, colored green for correct predictions and red for wrong ones.

The unnormalization step (images[i] * 0.5 + 0.5) reverses the normalization you applied when loading the data, converting pixel values from [-1, 1] back to [0, 1] so matplotlib can display them correctly.

torch.save(model.state_dict(), 'cifar10_model.pth') saves the trained weights to a file. You can reload this file later without retraining.

  • Save classifier.py.
  • Run the script by running this command:
python classifier.py

What should I see?

The script re-trains the model first (about 5-10 minutes on CPU), then prints the evaluation results. After training completes, you should see:

  • An overall test accuracy around 70-75%.
  • A per-class accuracy table where vehicles (ship, truck, automobile) tend to score higher than animals (cat, dog, deer).
  • predictions.png showing 16 test images with green (correct) and red (wrong) labels.
  • cifar10_model.pth saved to your project folder.

Not seeing the expected output?

  • If you see an indentation or syntax error, make sure the new code blocks are at the top level of the file (not indented inside the training loop or another function).
  • If predictions.png is not created, check that show_predictions(model, testloader) is called after the function definition (not inside it).
  • If accuracy is much lower than 70%, make sure your training loop ran for all 10 epochs before the evaluation code.

Why is my evaluation output missing or incorrect?

  • Confirm cifar10_model.pth exists in your cifar10-classifier folder by running:
ls

You should see cifar10_model.pth, predictions.png, sample_images.png, and training_loss.png listed alongside your classifier.py and data folder.

✔️ 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 classifier.py file at this point. Compare it against yours and fix any differences.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt


# --- Data Loading ---
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

trainset = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=transform
)
trainloader = DataLoader(trainset, batch_size=64, shuffle=True, num_workers=0)

testset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=transform
)
testloader = DataLoader(testset, batch_size=64, shuffle=False, num_workers=0)

classes = ('airplane', 'automobile', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

print(f'{len(trainset)} training images, {len(testset)} test images')


# --- Visualize Samples ---
def show_images(loader):
    dataiter = iter(loader)
    images, labels = next(dataiter)

    fig, axes = plt.subplots(1, 8, figsize=(12, 2))
    for i in range(8):
        img = images[i] * 0.5 + 0.5
        axes[i].imshow(img.permute(1, 2, 0).numpy())
        axes[i].set_title(classes[labels[i]], fontsize=8)
        axes[i].axis('off')
    plt.tight_layout()
    plt.savefig('sample_images.png', dpi=150)
    plt.show()


show_images(trainloader)


# --- CNN Architecture ---
class ImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(128 * 4 * 4, 256)
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')

model = ImageClassifier().to(device)
print(model)

total_params = sum(p.numel() for p in model.parameters())
print(f'Total parameters: {total_params:,}')


# --- Training ---
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

num_epochs = 10
train_losses = []

for epoch in range(num_epochs):
    running_loss = 0.0
    model.train()

    for inputs, labels in trainloader:
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    epoch_loss = running_loss / len(trainloader)
    train_losses.append(epoch_loss)
    print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {epoch_loss:.4f}')

print('Training complete!')


# --- Plot Training Loss ---
plt.figure(figsize=(8, 4))
plt.plot(range(1, num_epochs + 1), train_losses, marker='o')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training Loss Over Epochs')
plt.grid(True)
plt.savefig('training_loss.png', dpi=150)
plt.show()


# --- Evaluation ---
model.eval()
correct = 0
total = 0

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

print(f'Test Accuracy: {100 * correct / total:.2f}%')


# --- Per-Class Accuracy ---
class_correct = {c: 0 for c in classes}
class_total = {c: 0 for c in classes}

with torch.no_grad():
    for inputs, labels in testloader:
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)
        _, predicted = torch.max(outputs, 1)
        for i in range(labels.size(0)):
            label = classes[labels[i]]
            class_total[label] += 1
            if predicted[i] == labels[i]:
                class_correct[label] += 1

print('\nPer-Class Accuracy:')
for cls in classes:
    accuracy = 100 * class_correct[cls] / class_total[cls]
    print(f'  {cls:12s}: {accuracy:.1f}%')


# --- Visualize Predictions ---
def show_predictions(model, loader):
    dataiter = iter(loader)
    images, labels = next(dataiter)

    model.eval()
    with torch.no_grad():
        outputs = model(images.to(device))
        _, predicted = torch.max(outputs, 1)

    fig, axes = plt.subplots(2, 8, figsize=(14, 4))
    for i in range(16):
        row = i // 8
        col = i % 8
        img = images[i] * 0.5 + 0.5
        axes[row, col].imshow(img.permute(1, 2, 0).numpy())
        color = 'green' if predicted[i] == labels[i] else 'red'
        axes[row, col].set_title(classes[predicted[i]], fontsize=8, color=color)
        axes[row, col].axis('off')
    plt.suptitle('Green = Correct, Red = Wrong')
    plt.tight_layout()
    plt.savefig('predictions.png', dpi=150)
    plt.show()


show_predictions(model, testloader)


# --- Save Model ---
torch.save(model.state_dict(), 'cifar10_model.pth')
print('Model saved to cifar10_model.pth')

Secret mission

Supercharge Your CNN

Your CNN is working, but professional ML engineers have a few more tricks up their sleeve. Can you boost your model's accuracy by adding Batch Normalization, Dropout, and data augmentation to push past 80%?

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally, so there are no ongoing costs. Decide whether to keep your files on disk, or delete them entirely.

Resources you used:

  • The cifar10-classifier project folder on your Desktop, containing classifier.py, generated images (sample_images.png, training_loss.png, predictions.png), and the saved model (cifar10_model.pth).
  • The downloaded CIFAR-10 dataset inside the data/ subfolder (~170 MB).
  • The Python virtual environment inside the venv/ subfolder, with PyTorch and matplotlib installed.

Keep everything running

No action needed. Choose this if you want to keep experimenting with your CNN or try training with different hyperparameters.

  • Your project folder stays on your Desktop at cifar10-classifier/.
  • The CIFAR-10 dataset stays cached in data/ so future runs skip the download.
  • Your saved model weights (cifar10_model.pth) can be reloaded anytime without retraining.

Pause - I'll come back to this later

Deactivate the virtual environment to free your terminal session. Your files and installed packages stay on disk for when you return.

  • Deactivate the virtual environment by running:
deactivate
  • When you come back, navigate to cifar10-classifier/ and reactivate the virtual environment.

Delete - I don't want to use this again

Remove the entire project folder to delete all project files, the downloaded dataset, the virtual environment, and all installed packages.

  • If your virtual environment is still active, deactivate it first by running:
deactivate
  • Delete the entire project folder by running this command.

macOS

rm -rf ~/Desktop/cifar10-classifier

Windows

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

This removes classifier.py, all generated images, the saved model, the CIFAR-10 dataset, and the virtual environment in one step.

Nice Work!

Nice Work!

You just built a convolutional neural network from scratch and trained it to recognize objects in photographs. That is a foundational deep learning skill that every ML specialization builds on top of.

You've learned how to:

  • Build a CNN image classifier from scratch using PyTorch, with three convolutional blocks and two fully connected layers.
  • Write a complete training loop with forward pass, loss computation, backpropagation, and optimizer updates.
  • Evaluate model performance with per-class accuracy metrics and visualize predictions on test images.
  • Secret Mission: Supercharged your CNN with Batch Normalization, Dropout, data augmentation, and a learning rate scheduler to boost accuracy from ~72% to ~80%.

Ready to quiz yourself?