Build a CNN Image Classifier

Train a CNN on CIFAR-10 and optimize it from 55% to 83% accuracy.

Introduction

30 Second Summary

Your phone recognizes your face in a split second. Google Photos groups pictures of your dog together automatically. These features all rely on neural networks that have learned to "see" patterns in images.

In this project, you will build a convolutional neural network from scratch using PyTorch, train it on 60,000 real-world images, and progressively optimize it until it classifies objects with over 80% accuracy.

What You'll Build

You will run a notebook cell, watch your model correctly label a photo of a ship or a frog, then scroll up to a chart that proves each optimization technique you applied made the model measurably smarter.

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

  • A trained CNN that takes any 32x32 image and predicts which of 10 object categories it belongs to, from planes to trucks.
  • Training curves that visually expose overfitting in real time so you can see your model memorizing instead of learning.
  • A fully optimized model reaching ~83% accuracy after applying batch normalization, data augmentation, and learning rate scheduling.
  • Secret Mission: Discover which specific objects your model confuses most, then target those weaknesses with dropout regularization.

Are there any prerequisites?

You need basic Python knowledge (variables, functions, loops). No prior machine learning or deep learning experience is required.

Before We Start

Before we dive in, let's lock in what you're about to build and why it matters.

Set Up Your Deep Learning Environment

To build and train a convolutional neural network, you need PyTorch and its supporting libraries installed locally. PyTorch handles the math-heavy operations behind neural network training.

You'll install the CPU-only version of PyTorch in this step. Without the correct CPU-only packages, the training code could fail or try to find a GPU that is not available on your machine.

In this step, get ready to:
  • Install PyTorch (CPU-only), torchvision, and matplotlib in a virtual environment.
  • Create a Jupyter notebook in VS Code.
  • Verify all imports and confirm the compute device.
Create a virtual environment and install PyTorch

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

  • In VS Code, click Terminal in the top menu bar.
  • Click New Terminal to open a terminal panel at the bottom of the window.
  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir cnn-classifier
cd cnn-classifier
  • Create and activate a virtual environment by running these commands:

macOS

python3 -m venv venv
source venv/bin/activate

Windows

python -m venv venv
venv\Scripts\activate

Your terminal prompt now shows (venv) at the beginning, confirming the virtual environment is active.

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

Why CPU-only?

The --index-url flag points pip at PyTorch's CPU-only package index. This gives you a smaller download and avoids issues where code tries to find a GPU that does not exist on your machine.

CIFAR-10 images are only 32x32 pixels, so training a CNN on them takes minutes on CPU. No GPU needed for this project.

You will see pip downloading several packages. This may take a minute or two depending on your internet connection.

  • Install matplotlib (for plotting training curves) by running:
pip install matplotlib

Once both installs finish, you should see Successfully installed messages in your terminal output.

Install command failing?

  • Make sure you see (venv) in your prompt. If not, re-run the activation command above.
  • If you see a permissions error, do NOT use sudo. Re-activate your virtual environment instead.
  • If you see a network timeout, check your internet connection and try the command again.

Help me troubleshoot my PyTorch installation error.

Create your Jupyter notebook

Now that your packages are installed, you need to open this project folder in VS Code and create the notebook where you will write all your CNN code.

  • In VS Code, click File in the top menu bar.
  • Click Open Folder.
  • Navigate to your Desktop, select the cnn-classifier folder, and click Open.

What if VS Code asks to trust the folder?

VS Code may show a dialog asking if you trust the authors of the files in this folder. Click Yes, I trust the authors to continue.

  • In the Explorer sidebar on the left, hover over the folder name and click the New File icon (a page with a plus sign).
  • Type cnn_classifier.ipynb as the file name and press Enter.

VS Code recognizes the .ipynb extension and opens the file as a Jupyter notebook with an empty code cell ready for input.

  • Click Select Kernel in the top-right corner of the notebook.
  • Click Python Environments.
  • Select the venv Python interpreter from the list (it shows the path to your cnn-classifier/venv folder).

Why select a kernel?

The kernel tells VS Code which Python environment to run your notebook cells in. By selecting the venv kernel, your notebook uses the PyTorch and matplotlib packages you just installed.

Don't see your venv in the kernel list?

  • Press Cmd+Shift+P (macOS) or Ctrl+Shift+P (Windows) and type Python: Select Interpreter. Manually browse to the venv Python inside your cnn-classifier/venv folder.
  • If you still do not see it, try closing and reopening VS Code with the cnn-classifier folder open.

Help me find my virtual environment kernel in VS Code.

Import libraries and confirm your setup

Your notebook is connected to the virtual environment. Now you will write the first cell that imports every library the project needs and confirms PyTorch is running on CPU.

  • Click inside the empty code cell in your notebook and type the following code:
# Cell 1: Imports and Setup
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import time

device = torch.device('cpu')
print(f'PyTorch version: {torch.__version__}')
print(f'Using device: {device}')

What does this code do?

  • torch is the core PyTorch library. torch.nn provides neural network layers, and torch.optim provides optimizers that update the network's weights during training.
  • torchvision handles image datasets and transforms. You will use it to download CIFAR-10 in the next step.
  • matplotlib.pyplot plots training curves so you can visualize how your model improves over time.
  • device = torch.device('cpu') explicitly sets the compute device to CPU. This ensures all tensor operations run on your processor.

Before you run this cell, what do you expect to see printed? Think about what version of PyTorch you just installed and which device you set.

  • Run the cell by pressing Shift+Enter or clicking the play button to the left of the cell.

You should see output like this:

PyTorch version: 2.13.0
Using device: cpu

What should I see?

The version number confirms PyTorch installed correctly. Your exact version number may differ slightly, but it should start with 2. and the device must say cpu.

Seeing a ModuleNotFoundError?

  • This means your notebook kernel is not pointing at the virtual environment. Click the kernel name in the top-right corner and re-select the venv Python interpreter.
  • If the error mentions ipykernel, VS Code may prompt you to install it. Click Install when prompted.

Help me fix my ModuleNotFoundError when running my notebook.

✔️ Awesome, I've got everything!

Great. Make sure you have saved your notebook with Cmd+S (macOS) or Ctrl+S (Windows).

ⓧ I'd like to double check the full code

Here is the complete Cell 1. Make sure your notebook cell matches this exactly:

# Cell 1: Imports and Setup
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import time

device = torch.device('cpu')
print(f'PyTorch version: {torch.__version__}')
print(f'Using device: {device}')

Your deep learning environment is ready. Next up, you will download 60,000 real-world images and train your first neural network on them.

Load CIFAR-10 and Train a Basic CNN

Before you can improve a neural network, you need a baseline to measure against. A simple CNN trained on real images gives you a concrete accuracy number to beat in later steps.

In this step, you'll load 50,000 real-world images from the CIFAR-10 dataset, build a 3-layer CNN from scratch, train it for 10 epochs, and plot the results to see exactly where it falls short.

In this step, get ready to:
  • Download CIFAR-10 and display sample images to confirm the data loaded correctly.
  • Define a BasicCNN class with convolutional layers, pooling, and fully-connected layers.
  • Train the model for 10 epochs and plot accuracy curves that reveal overfitting.
Download CIFAR-10 and visualize sample images

CIFAR-10 contains 60,000 tiny color images (32x32 pixels) across 10 categories like planes, cars, and cats. You'll download it through torchvision and wrap it in DataLoaders that feed batches of 64 images to your model during training.

  • In your notebook, create a new cell below Cell 1 and paste the following data loading code:
# Cell 2: Load CIFAR-10 with basic transforms
transform_basic = 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_basic)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)

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

classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

What does this code do?

  • transform_basic converts each image to a tensor and normalizes pixel values from [0, 1] to [-1, 1]. This centering helps the network learn faster.
  • torchvision.datasets.CIFAR10 downloads the dataset (~163 MB) to a ./data folder the first time you run it.
  • DataLoader wraps the dataset so your model receives shuffled batches of 64 images during training, rather than one image at a time.
  • classes maps the numeric labels (0-9) to human-readable category names.
  • In the same cell, add the visualization code below the classes line:
# Display sample images
dataiter = iter(trainloader)
images, labels = next(dataiter)

fig, axes = plt.subplots(2, 5, figsize=(12, 5))
for i, ax in enumerate(axes.flat):
    img = images[i] * 0.5 + 0.5  # unnormalize
    ax.imshow(img.permute(1, 2, 0).numpy())
    ax.set_title(classes[labels[i]])
    ax.axis('off')
plt.suptitle('CIFAR-10 Sample Images')
plt.tight_layout()
plt.show()

print(f'Training set: {len(trainset)} images')
print(f'Test set: {len(testset)} images')
print(f'Image shape: {images[0].shape}')
print(f'Classes: {classes}')

What does the visualization do?

This grabs 10 images from the first batch, reverses the normalization (multiplying by 0.5 and adding 0.5 to get back to [0, 1]), and displays them in a 2x5 grid with their class labels. The print statements confirm dataset sizes and image dimensions.

  • Run Cell 2.

You should see a grid of 10 small color images with labels like frog, truck, or ship above each one. The output also confirms 50,000 training images and 10,000 test images.

Seeing a download error?

  • Make sure you have an internet connection. The first run downloads CIFAR-10 from the internet.
  • If the download stalls, interrupt the cell and run it again. torchvision will resume from the cached file.
  • help me debug a CIFAR-10 download error in my Jupyter notebook
Define the BasicCNN architecture

A Convolutional Neural Network applies learnable filters to images, detecting increasingly complex patterns at each layer. Your first CNN uses three convolutional layers that shrink the image from 32x32 down to 4x4, then flattens the result into a prediction across 10 classes.

  • Create a new cell (Cell 3) and paste the model definition:
# Cell 3: Define Basic CNN
class BasicCNN(nn.Module):
    def __init__(self):
        super(BasicCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 64, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 4 * 4, 512)
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))   # 32x32 -> 16x16
        x = self.pool(torch.relu(self.conv2(x)))   # 16x16 -> 8x8
        x = self.pool(torch.relu(self.conv3(x)))   # 8x8 -> 4x4
        x = x.view(-1, 64 * 4 * 4)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model_v1 = BasicCNN()
total_params = sum(p.numel() for p in model_v1.parameters())
print(f'BasicCNN architecture:')
print(model_v1)
print(f'\nTotal parameters: {total_params:,}')

What does each layer do?

  • conv1 takes 3 input channels (RGB) and outputs 32 feature maps. Each subsequent conv layer detects more complex patterns by combining the previous layer's features.
  • MaxPool2d(2, 2) halves the spatial dimensions after each conv layer, reducing 32x32 to 16x16 to 8x8 to 4x4.
  • fc1 flattens the 4x4x64 feature maps into 512 neurons. fc2 outputs 10 scores, one per class.
  • Run Cell 3.

You should see the full model architecture printed, followed by a total parameter count in the hundreds of thousands.

Seeing a NameError?

  • Make sure Cell 1 ran successfully first. The nn and torch imports must be loaded before defining the model.
  • Check that the class definition starts with class BasicCNN(nn.Module): with the correct capitalization.
  • help me fix a NameError when defining my BasicCNN class
Write a training loop and evaluate the model

A training loop is the engine that teaches your network. Each epoch passes all 50,000 images through the model, computes how wrong each prediction is, and adjusts the model's weights to reduce that error. You'll build a reusable function that tracks loss and accuracy on both training and test data.

  • Create a new cell (Cell 4) and start writing the training function with the forward pass and weight update logic:
# Cell 4: Training function
def train_model(model, trainloader, testloader, criterion, optimizer, epochs, scheduler=None):
    train_losses = []
    train_accuracies = []
    test_accuracies = []

    for epoch in range(epochs):
        model.train()
        running_loss = 0.0
        correct_train = 0
        total_train = 0

        for images, labels in trainloader:
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            _, predicted = torch.max(outputs, 1)
            total_train += labels.size(0)
            correct_train += (predicted == labels).sum().item()

        if scheduler:
            scheduler.step()

        train_loss = running_loss / len(trainloader)
        train_acc = 100 * correct_train / total_train
        train_losses.append(train_loss)
        train_accuracies.append(train_acc)

What does the training loop do?

  • optimizer.zero_grad() clears old gradients so each batch starts fresh.
  • loss.backward() computes how much each weight contributed to the error. optimizer.step() nudges the weights in the direction that reduces loss.
  • The loop accumulates running loss and correct predictions to compute per-epoch averages.
  • In the same cell, continue the function directly below with the evaluation section:
        model.eval()
        correct_test = 0
        total_test = 0
        with torch.no_grad():
            for images, labels in testloader:
                outputs = model(images)
                _, predicted = torch.max(outputs, 1)
                total_test += labels.size(0)
                correct_test += (predicted == labels).sum().item()

        test_acc = 100 * correct_test / total_test
        test_accuracies.append(test_acc)

        lr_str = f' | LR: {scheduler.get_last_lr()[0]:.6f}' if scheduler else ''
        print(f'Epoch {epoch+1}/{epochs} | Loss: {train_loss:.3f} | Train Acc: {train_acc:.1f}% | Test Acc: {test_acc:.1f}%{lr_str}')

    return train_losses, train_accuracies, test_accuracies

What does the evaluation section do?

  • model.eval() switches from training mode to evaluation mode, disabling randomness like dropout.
  • torch.no_grad() tells PyTorch not to track gradients during testing. This saves memory and speeds up evaluation.
  • After each epoch, the function prints a progress line showing loss, train accuracy, and test accuracy side by side.
  • Run Cell 4 to define the function.

No output appears yet. The function is now defined and ready to be called.

Now it's time to train. Before you run this next cell, think about what accuracy you expect from a model that has never seen these images before. Random guessing across 10 classes would give 10%.

  • Create a new cell (Cell 5) and kick off training by pasting:
# Cell 5: Train the basic model
criterion = nn.CrossEntropyLoss()
optimizer_v1 = optim.Adam(model_v1.parameters(), lr=0.001)

print('Training BasicCNN for 10 epochs...')
start_time = time.time()
losses_v1, train_accs_v1, test_accs_v1 = train_model(model_v1, trainloader, testloader, criterion, optimizer_v1, epochs=10)
elapsed = time.time() - start_time
print(f'\nTraining completed in {elapsed:.1f} seconds')
print(f'Final test accuracy: {test_accs_v1[-1]:.1f}%')

What does this code do?

  • CrossEntropyLoss is the standard loss function for multi-class classification. It measures how far off each prediction is from the true label.
  • Adam is an optimizer that adapts the learning rate for each weight individually. lr=0.001 sets how big each weight update step is.
  • Training on CPU will take approximately 3-5 minutes for 10 epochs. You'll see progress printed after each epoch.
  • Run Cell 5 and wait for all 10 epochs to complete.

You should see epoch-by-epoch output with training accuracy climbing toward ~75% but test accuracy plateauing around 55-62%. Notice how training accuracy keeps rising while test accuracy stalls. That gap is the problem.

Do you think the model would keep improving on test data if you trained it for more epochs? Before you plot the curves, predict what the graph will look like.

  • Create a new cell (Cell 6) and paste the plotting code to visualize the training history:
# Cell 6: Plot basic model results
fig, axes = plt.subplots(1, 2, figsize=(12, 4))

axes[0].plot(losses_v1)
axes[0].set_title('Training Loss (BasicCNN)')
axes[0].set_xlabel('Epoch')
axes[0].set_ylabel('Loss')

axes[1].plot(train_accs_v1, label='Train Accuracy')
axes[1].plot(test_accs_v1, label='Test Accuracy')
axes[1].set_title('Accuracy (BasicCNN)')
axes[1].set_xlabel('Epoch')
axes[1].set_ylabel('Accuracy (%)')
axes[1].legend()
axes[1].set_ylim([0, 100])

plt.tight_layout()
plt.show()

print(f'Gap between train and test accuracy: {train_accs_v1[-1] - test_accs_v1[-1]:.1f}%')
print('The model is OVERFITTING - it memorizes training data but fails on new images.')

What do these plots show?

The left plot shows training loss decreasing over epochs. A falling loss means the model is learning to make better predictions on training data.

The right plot is the crucial one. Two lines diverge: train accuracy climbs to ~75% while test accuracy plateaus at ~55-62%. This visible gap is called overfitting. The model memorizes training images instead of learning general features that transfer to new images.

  • Run Cell 6.

You should see two side-by-side plots. The accuracy plot shows a clear gap between the train and test lines.

Seeing a NameError on losses_v1?

  • Make sure Cell 5 ran to completion. If you interrupted training early, the variables won't exist.
  • Re-run Cell 5 and let all 10 epochs finish before running Cell 6.
  • help me fix a NameError when plotting my training results

✔️ Awesome, I've got everything!

Great. Double check you've saved your notebook with all 6 cells.

ⓧ I'd like to double check the full code

Your notebook should contain these cells in order. Scroll up and confirm each one is present:

  • Cell 1: Imports and Setup (from Step 1).
  • Cell 2: Load CIFAR-10 with basic transforms (data loading, DataLoaders, classes tuple, image grid display, and print statements all in one cell).
  • Cell 3: Define Basic CNN (BasicCNN class with __init__ and forward methods, then model instantiation and parameter count).
  • Cell 4: Training function (the full train_model function with training loop and evaluation in one cell).
  • Cell 5: Train the basic model (criterion, optimizer, and the train_model call for 10 epochs).
  • Cell 6: Plot basic model results (two subplots for loss and accuracy, plus the overfitting gap print).

Your baseline is set. The model learns the training data well but chokes on new images. In the next step, you'll apply batch normalization to close that gap and push test accuracy significantly higher.

Fix Overfitting with Batch Normalization

Your BasicCNN is trained and the accuracy plots tell a clear story. The model scores around 75% on training images but drops to 55-62% on new test images it has never seen before.

That gap is overfitting. The model memorized specific pixel patterns in the training set instead of learning general features like "what makes a truck look like a truck." Batch normalization fixes this by normalizing the inputs at each layer, which stabilizes training and forces the network to learn more robust features.

In this step, get ready to:
  • Define an ImprovedCNN class with BatchNorm2d layers after each convolution.
  • Train the improved model and plot a side-by-side comparison with the basic model.
Define the ImprovedCNN class with batch normalization

Batch normalization normalizes the output of each convolutional layer so the next layer receives a consistent input distribution. You will also increase the third conv layer from 64 to 128 filters to give the model more learning capacity.

What is batch normalization?

During training, the output distribution of each layer shifts as weights update. This makes later layers constantly chase a moving target.

BatchNorm2d normalizes each mini-batch to have mean 0 and variance 1 before passing data to the next layer. This keeps the input distribution stable so the network trains faster and generalizes better.

  • Create a new cell in your notebook (Cell 7) and define the ImprovedCNN class by pasting this code:
# Cell 7: Define Improved CNN with Batch Normalization
class ImprovedCNN(nn.Module):
    def __init__(self):
        super(ImprovedCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(128 * 4 * 4, 512)
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.bn1(self.conv1(x))))
        x = self.pool(torch.relu(self.bn2(self.conv2(x))))
        x = self.pool(torch.relu(self.bn3(self.conv3(x))))
        x = x.view(-1, 128 * 4 * 4)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model_v2 = ImprovedCNN()
total_params_v2 = sum(p.numel() for p in model_v2.parameters())
print(f'ImprovedCNN with BatchNorm:')
print(model_v2)
print(f'\nTotal parameters: {total_params_v2:,}')

What changed from BasicCNN?

  • Three BatchNorm2d layers are inserted after each convolution and before the activation. The ordering is conv -> bn -> relu -> pool.
  • The third conv layer now outputs 128 channels instead of 64. This gives the deeper layers more capacity to represent complex features.
  • The fully-connected layer input changes from 64 * 4 * 4 to 128 * 4 * 4 to match the new channel count.
  • Run the cell.
  • You should see the model architecture printed with BatchNorm2d layers listed after each Conv2d, and a total parameter count higher than the BasicCNN.

Seeing a shape mismatch error?

Check that self.fc1 uses 128 * 4 * 4 (not 64 * 4 * 4). The third conv layer now outputs 128 channels, so the flatten size changed.

Also verify that x.view(-1, 128 * 4 * 4) matches the fc1 input dimension.

Help me fix this shape mismatch in my ImprovedCNN

Train the improved model and compare results

Now you will train this improved model using the exact same data, loss function, and optimizer settings. The only difference is the architecture. This makes it a fair comparison.

Before you run training, what do you think will happen to the gap between train and test accuracy? Will batch normalization narrow it, widen it, or leave it the same?

  • Create a new cell (Cell 8) and train the improved model by running this code:
# Cell 8: Train improved model
optimizer_v2 = optim.Adam(model_v2.parameters(), lr=0.001)

print('Training ImprovedCNN (with BatchNorm) for 10 epochs...')
start_time = time.time()
losses_v2, train_accs_v2, test_accs_v2 = train_model(model_v2, trainloader, testloader, criterion, optimizer_v2, epochs=10)
elapsed = time.time() - start_time
print(f'\nTraining completed in {elapsed:.1f} seconds')
print(f'Final test accuracy: {test_accs_v2[-1]:.1f}%')

What does this code do?

  • optimizer_v2 creates a fresh Adam optimizer for the new model with the same learning rate (0.001) as before.
  • The train_model function you defined in Step 2 handles the full training loop. Reusing it keeps the comparison fair.
  • The results are stored in losses_v2, train_accs_v2, and test_accs_v2 so you can compare them against the basic model later.
  • Run the cell. Training will take a few minutes on CPU.
  • Watch the epoch-by-epoch output. You should see the final test accuracy printed at around 70-75%, a significant jump from the BasicCNN's 55-62%.

Training seems stuck at low accuracy?

Make sure you ran Cell 7 first to define and instantiate model_v2. If you re-run Cell 8 without re-running Cell 7, you are training an already-partially-trained model with a reset optimizer, which produces unpredictable results.

Help me debug my ImprovedCNN training accuracy

Now let's visualize the improvement by plotting both models on the same chart.

  • Create a new cell (Cell 9) and generate the comparison plot by running this code:
# Cell 9: Compare basic vs improved
fig, ax = plt.subplots(1, 1, figsize=(8, 5))
ax.plot(test_accs_v1, label='BasicCNN (no BatchNorm)', linestyle='--')
ax.plot(test_accs_v2, label='ImprovedCNN (with BatchNorm)')
ax.set_title('Test Accuracy Comparison')
ax.set_xlabel('Epoch')
ax.set_ylabel('Test Accuracy (%)')
ax.legend()
ax.set_ylim([40, 85])
plt.tight_layout()
plt.show()

improvement = test_accs_v2[-1] - test_accs_v1[-1]
print(f'Improvement from BatchNorm: +{improvement:.1f}%')
print(f'Train/test gap reduced from {train_accs_v1[-1] - test_accs_v1[-1]:.1f}% to {train_accs_v2[-1] - test_accs_v2[-1]:.1f}%')

What does this plot show?

The dashed line is your BasicCNN test accuracy from Step 2. The solid line is the ImprovedCNN with batch normalization. The vertical gap between them is the accuracy gained from a single architectural change.

The printed gap reduction confirms the overfitting is shrinking. A smaller train/test gap means the model is learning general features rather than memorizing training-specific patterns.

  • Run the cell.
  • You should see a line chart with the ImprovedCNN line clearly above the BasicCNN line, and a printed message showing the improvement (around +10-15%) and a reduced train/test gap.

Plot shows both lines at the same level?

Make sure Cell 5 (BasicCNN training) was run in a previous session and test_accs_v1 still holds data. If you restarted the kernel, you need to re-run Cells 1 through 6 before running Cells 7-9.

Help me fix my comparison plot

✔️ Awesome, I've got everything!

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

ⓧ I'd like to double check the full code

Scroll through your notebook and confirm these cells are present in order:

  • Cell 1: Imports and Setup (from Step 1).
  • Cell 2: Load CIFAR-10 with basic transforms (from Step 2).
  • Cell 3: Define Basic CNN (from Step 2).
  • Cell 4: Training function (from Step 2).
  • Cell 5: Train the basic model (from Step 2).
  • Cell 6: Plot basic model results (from Step 2).
  • Cell 7: Define Improved CNN with Batch Normalization (this step).
  • Cell 8: Train improved model (this step).
  • Cell 9: Compare basic vs improved (this step).

Batch normalization took your test accuracy from the mid-50s into the low-70s with just one architectural change. But there is still room to improve. Next up, you will teach the model to generalize even further using data augmentation and a learning rate scheduler.

Boost Generalization with Data Augmentation and LR Scheduling

Batch normalization stabilized your training and narrowed the accuracy gap. But the model still only sees each training image once per epoch in its original form.

In this step, you'll apply data augmentation to create flipped and cropped variations of every training image. Then you'll add a learning rate scheduler that gradually reduces the learning rate so the model can fine-tune in later epochs without overshooting.

In this step, get ready to:
  • Create an augmented transform pipeline and visualize its effect on training images.
  • Train the ImprovedCNN with augmentation and a StepLR scheduler for 20 epochs.
  • Compare all three model versions in a final bar chart showing progressive improvement.
Build an augmented transform pipeline

Your current transform_basic pipeline only converts images to tensors and normalizes them. Every epoch, the model sees the exact same 50,000 images. Data augmentation randomly flips and crops each image on the fly, so the model effectively trains on a much larger variety of examples.

  • Create a new cell below Cell 9 and define the augmented transform pipeline with its DataLoader by pasting this code:
# Cell 10: Data augmentation transforms
transform_augmented = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32, padding=4),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

trainset_aug = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_augmented)
trainloader_aug = torch.utils.data.DataLoader(trainset_aug, batch_size=64, shuffle=True)

What does this code do?

RandomHorizontalFlip() randomly mirrors images left-to-right with 50% probability. A flipped car is still a car, so this teaches the model that orientation doesn't change the label.

RandomCrop(32, padding=4) pads 4 pixels of black around the image, then randomly crops back to 32x32. This forces the model to recognize objects even when they're slightly off-center.

This cell won't produce any visual output yet. The next chunk visualizes the difference between original and augmented images.

  • In the same cell, add the visualization code directly below what you just pasted:
# Show augmented vs original
fig, axes = plt.subplots(2, 5, figsize=(12, 5))
fig.suptitle('Top: Original | Bottom: Augmented')

orig_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, transform=transforms.ToTensor())
for i in range(5):
    img, label = orig_dataset[i]
    axes[0, i].imshow(img.permute(1, 2, 0).numpy())
    axes[0, i].set_title(classes[label])
    axes[0, i].axis('off')

aug_viz = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32, padding=4),
    transforms.ToTensor()
])
aug_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, transform=aug_viz)
for i in range(5):
    img, label = aug_dataset[i]
    axes[1, i].imshow(img.permute(1, 2, 0).numpy())
    axes[1, i].set_title(classes[label])
    axes[1, i].axis('off')

plt.tight_layout()
plt.show()

What does this visualization show?

The top row shows the original training images. The bottom row shows the same images after random augmentation. You'll see some images flipped horizontally and others slightly shifted from the random crop.

Each time you run this cell, the augmented row looks different. That's the point. The model never sees the exact same image twice.

  • Run Cell 10.
  • You should see a 2x5 grid of images. The top row shows original images and the bottom row shows augmented versions with visible flips or slight position shifts.

Not seeing the image grid?

Make sure you ran Cells 1 and 2 first. The classes variable must already exist in your notebook session.

If you see a NameError, restart your kernel and re-run all cells from the top.

Help me debug my augmentation visualization.

Train with augmentation and a learning rate scheduler

Now you'll combine data augmentation with a learning rate scheduler. The StepLR scheduler multiplies the learning rate by a factor (gamma) every N epochs. This lets the optimizer take large steps early in training and smaller, more precise steps later.

  • Create a new cell below Cell 10 and set up the model, optimizer, and scheduler by pasting this code:
# Cell 11: Train with augmentation and LR scheduler
model_v3 = ImprovedCNN()
optimizer_v3 = optim.Adam(model_v3.parameters(), lr=0.001)
scheduler = optim.lr_scheduler.StepLR(optimizer_v3, step_size=10, gamma=0.1)

print('Training ImprovedCNN with augmentation + LR scheduler for 20 epochs...')
start_time = time.time()
losses_v3, train_accs_v3, test_accs_v3 = train_model(
    model_v3, trainloader_aug, testloader, criterion, optimizer_v3, epochs=20, scheduler=scheduler
)
elapsed = time.time() - start_time
print(f'\nTraining completed in {elapsed:.1f} seconds')
print(f'Final test accuracy: {test_accs_v3[-1]:.1f}%')

What does this code do?

  • model_v3 = ImprovedCNN() creates a fresh copy of the batch-normalized architecture. You train from scratch so results are directly comparable.
  • StepLR(optimizer_v3, step_size=10, gamma=0.1) multiplies the learning rate by 0.1 every 10 epochs. After epoch 10, the LR drops from 0.001 to 0.0001.
  • The model trains on trainloader_aug (augmented data) but evaluates on the original testloader (unchanged test images).

Before you run this cell, what test accuracy do you predict? The batch-norm model reached ~72%. How much do you think augmentation and scheduling will add?

  • Run Cell 11. This will take approximately 6-10 minutes on CPU since it trains for 20 epochs.

You should see 20 lines of epoch output. Notice the LR value printed after epoch 10 drops from 0.001000 to 0.000100. The final test accuracy should land around 82-85%.

Why does accuracy jump after epoch 10?

When the scheduler drops the learning rate at epoch 10, the optimizer stops overshooting and starts settling into a sharper minimum. You'll often see a noticeable accuracy bump right after the LR drop.

Accuracy stuck below 75%?

Check that you're passing trainloader_aug (not trainloader) as the second argument to train_model. Without augmentation, you're just re-training the same model as Step 3.

Also confirm scheduler=scheduler is passed to the function. Without it, the learning rate never drops.

Help me figure out why my augmented training isn't improving accuracy.

Compare all three models side by side

You now have three sets of results stored in your notebook. A bar chart makes the progressive improvement obvious at a glance.

  • Create a new cell below Cell 11 and generate the final comparison chart by pasting this code:
# Cell 12: Final comparison of all three models
final_accuracies = [test_accs_v1[-1], test_accs_v2[-1], test_accs_v3[-1]]
model_names = ['BasicCNN\n(baseline)', 'BatchNorm\n(10 epochs)', 'BatchNorm +\nAugmentation +\nLR Schedule\n(20 epochs)']

fig, ax = plt.subplots(figsize=(8, 5))
bars = ax.bar(model_names, final_accuracies, color=['#ff6b6b', '#feca57', '#48dbfb'])
ax.set_ylabel('Test Accuracy (%)')
ax.set_title('Progressive CNN Improvement on CIFAR-10')
ax.set_ylim([0, 100])

for bar, acc in zip(bars, final_accuracies):
    ax.text(bar.get_x() + bar.get_width()/2., bar.get_height() + 1, f'{acc:.1f}%',
            ha='center', va='bottom', fontweight='bold')

plt.tight_layout()
plt.show()

print('\nSummary of techniques and impact:')
print(f'  Baseline (3-layer CNN):                    {final_accuracies[0]:.1f}%')
print(f'  + Batch Normalization:                     {final_accuracies[1]:.1f}% (+{final_accuracies[1]-final_accuracies[0]:.1f}%)')
print(f'  + Data Augmentation + LR Scheduling:       {final_accuracies[2]:.1f}% (+{final_accuracies[2]-final_accuracies[0]:.1f}% total)')

What does this code do?

It pulls the final test accuracy from each model's history (test_accs_v1, test_accs_v2, test_accs_v3) and plots them as a color-coded bar chart. The text annotations above each bar show the exact percentage.

  • Run Cell 12.
  • You should see a bar chart with three bars climbing from left to right: BasicCNN around 58%, BatchNorm around 72%, and the fully optimized model around 83%.
  • Below the chart, a printed summary shows the exact improvement from each technique.

Seeing a NameError for one of the accuracy lists?

All three training results (test_accs_v1, test_accs_v2, test_accs_v3) must exist in memory. If you restarted your kernel at any point, re-run Cells 5, 8, and 11 to regenerate the data.

Help me fix my comparison chart cell.

✔️ Awesome, I've got everything!

Great. Make sure your notebook is saved. You should have Cells 1-12 running cleanly with the final bar chart visible.

ⓧ I'd like to double check the full code

Your notebook now has 12 cells. Scroll through and confirm each one is present in this order:

  • Cell 1: Imports and Setup (torch, torchvision, matplotlib, numpy, time).
  • Cell 2: Load CIFAR-10 with basic transforms, display sample grid.
  • Cell 3: Define BasicCNN class.
  • Cell 4: train_model function definition.
  • Cell 5: Train BasicCNN for 10 epochs.
  • Cell 6: Plot basic model loss and accuracy curves.
  • Cell 7: Define ImprovedCNN with BatchNorm.
  • Cell 8: Train ImprovedCNN for 10 epochs.
  • Cell 9: Compare BasicCNN vs ImprovedCNN test accuracy.
  • Cell 10: Augmented transforms, DataLoader, and augmentation visualization.
  • Cell 11: Train with augmentation + StepLR scheduler for 20 epochs.
  • Cell 12: Final bar chart comparing all three models.

You took a model that barely beat random guessing and pushed it above 80% accuracy using just three techniques. Your notebook now tells the full visual story of how each optimization contributes to a better model. Next up: a secret mission that digs into which specific classes your model still struggles with.

Secret mission

Build a Confusion Matrix and Add Dropout

Your model hits ~83% overall accuracy, but which specific objects is it getting wrong? Some classes might be at 90% while others languish at 65%. In this mission, you will build a confusion matrix to expose exactly where your model struggles, then apply dropout regularization to improve its weakest predictions.

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 for future experiments or delete them to free up disk space.

Resources you used:

  • The ./data folder containing the CIFAR-10 dataset cache (~170 MB).
  • Your cnn_classifier.ipynb notebook with all model definitions, training history, and plots.
  • A Python virtual environment with PyTorch, torchvision, and matplotlib installed.

Keep everything running

No action needed. Choose this if you plan to continue experimenting with CNNs or want to reuse the CIFAR-10 dataset for other PyTorch projects.

  • The ./data folder stays cached so future projects can load CIFAR-10 without re-downloading.
  • Your virtual environment keeps PyTorch and all dependencies ready to use.
  • Your notebook preserves all training curves and accuracy comparisons for reference.

Pause - I'll come back to this later

Close your notebook and deactivate the virtual environment to free up memory. Your files remain on disk for when you return.

  • Save and close cnn_classifier.ipynb in VS Code.
  • In the terminal, deactivate your virtual environment by running:
deactivate
  • When you come back, reactivate the environment and reopen the notebook to pick up where you left off.

Delete - I don't want to use this again

Remove all project files, the dataset cache, and the virtual environment to reclaim disk space.

  • Deactivate the virtual environment first by running:
deactivate
  • Navigate to the parent directory of your project folder, then delete the entire project folder.

On macOS/Linux:

cd ~/Desktop
rm -rf cnn-classifier

On Windows:

cd ~\Desktop
rmdir /s /q cnn-classifier
  • This removes the notebook, the ./data folder (~170 MB), and the virtual environment in one step.

Export before deleting

If you want to keep a record of your training curves, export your notebook as HTML first. In VS Code, open the command palette and select Export Notebook As then choose HTML. Save it somewhere outside the project folder.

Nice Work!

Nice Work!

You just built, trained, and optimized a convolutional neural network from scratch using PyTorch. Your notebook tells the full story of taking a model from barely better than random guessing to over 80% accuracy on real-world images.

You've learned how to:

  • Built a CNN from scratch with Conv2d, MaxPool2d, and Linear layers, trained it on 50,000 CIFAR-10 images, and diagnosed overfitting by reading train vs test accuracy curves.
  • Applied batch normalization to stabilize training and narrow the accuracy gap, boosting test performance from ~58% to ~72%.
  • Used data augmentation and a learning rate scheduler to push the final model to ~83% accuracy, proving that each optimization technique contributes a measurable improvement.
  • Secret Mission: Built a confusion matrix heatmap to identify the weakest classes, then applied dropout regularization to reduce remaining overfitting and improve per-class accuracy.

Ready to quiz yourself?