CIFAR-10 Image Classifier with CNNs

Train three CNN models on CIFAR-10, reaching 90%+ accuracy with transfer learning.

Introduction

30 Second Summary

Every time you unlock your phone with your face or scroll through auto-tagged photos, an image classifier is doing the work behind the scenes. These models look at raw pixels and decide what they see.

In this project, you will build an image classification pipeline in PyTorch that learns to recognize 10 categories of real-world images. You will train three progressively better models and apply transfer learning to achieve 90%+ accuracy, all inside a Google Colab notebook.

What You'll Build

You'll run the final cell in your notebook and see a confusion matrix heatmap proving your model correctly classifies over 90% of test images, alongside a prediction grid where correct answers glow green and mistakes show red.

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

  • A simple CNN trained from scratch on CIFAR-10 that hits ~65% accuracy, giving you a baseline to measure every improvement against.
  • An improved CNN with data augmentation and dropout that pushes to ~73%, with training curve plots that visually expose the overfitting gap and prove each fix works.
  • A pretrained ResNet18 model that reaches 90%+ accuracy, evaluated with a confusion matrix heatmap, per-class accuracy breakdown, and a color-coded prediction grid showing correct and incorrect classifications.
  • Secret Mission: Add Grad-CAM heatmap visualization to see exactly which pixels your model focuses on when making predictions.

Are there any prerequisites?

You should be comfortable with Python basics. No prior machine learning experience is required. Everything runs in your browser through Google Colab, so there is nothing to install on your computer.

Before We Start

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

Set Up Your Colab Notebook

Training neural networks requires serious computing power. A GPU processes thousands of image calculations in parallel, turning hours of training time into minutes.

Google Colab gives you a free GPU-powered Python environment right in your browser. In this step, you'll create a Colab notebook, enable GPU acceleration, and import all the libraries your image classifier needs.

In this step, get ready to:
  • Open Google Colab and enable GPU acceleration.
  • Import PyTorch, torchvision, and supporting libraries.
  • Confirm your GPU is active and ready to train.
Open Google Colab and enable GPU

Google Colab is a free, browser-based notebook environment that comes with PyTorch preinstalled. You just need to create a notebook and switch on GPU acceleration.

  • Go to colab.research.google.com in your browser.
  • Sign in with your Google account if prompted.
  • Click New notebook to create a blank notebook.

By default, Colab runs on a CPU. You need to switch to a GPU runtime so your neural network training finishes in minutes instead of hours.

  • Click Runtime in the top menu bar.
  • Click Change runtime type.
  • Set Hardware accelerator to T4 GPU.
  • Click Save.

Why enable GPU?

Convolutional Neural Networks involve millions of matrix multiplications during training. A GPU handles thousands of these operations simultaneously, making training 10-50x faster than a CPU.

Don't see T4 GPU as an option?

Colab's free GPUs are shared across all users. If T4 GPU is unavailable, try selecting a different GPU type from the dropdown. If no GPU options appear, wait a few minutes and try again.

Help me set up GPU acceleration in Google Colab.

Import your libraries and verify GPU access

Your notebook needs several libraries to build and train a neural network. Colab comes with all of them preinstalled, so you just need to import them.

  • Click into the first code cell in your notebook and paste the following code:
# ============================================================
# STEP 1: Setup
# ============================================================
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torchvision.models import resnet18, ResNet18_Weights
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import numpy as np

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')
if device.type == 'cuda':
    print(f'GPU: {torch.cuda.get_device_name(0)}')
print(f'PyTorch version: {torch.__version__}')

What does this code do?

  • torch, torch.nn, and torch.optim are PyTorch's core modules for building and training neural networks.
  • torchvision provides image datasets, pretrained models, and image transformation tools.
  • matplotlib and numpy handle plotting training curves and working with arrays.
  • device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') checks for a GPU. If one is available, all training runs on the GPU instead of the slower CPU.
  • The three print statements confirm which device is active, the GPU name, and which PyTorch version you're running.

Before you run this cell, what do you think the output will say? Will the device be cuda (GPU) or cpu?

  • Press Shift+Enter to run the cell.

Since you enabled the T4 GPU, you should see three lines of output confirming everything is working.

What should I see?

Your cell output shows three lines. Using device: cuda confirms the GPU is active. GPU: Tesla T4 names the specific GPU. The third line prints your PyTorch version number.

Seeing 'Using device: cpu' instead?

  • Go back to Runtime > Change runtime type and confirm T4 GPU is selected. After changing the runtime, re-run your cell.
  • If no GPU options are available, Colab's free GPUs may be at capacity. Wait a few minutes and try switching the runtime type again.

Why is my Colab notebook not detecting a GPU?

✔️ Awesome, I've got everything!

Great. Make sure your notebook is saved by pressing Cmd+S (macOS) or Ctrl+S (Windows).

ⓧ I'd like to double check the full code

Here is the complete cell. Make sure your code matches this exactly:

# ============================================================
# STEP 1: Setup
# ============================================================
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torchvision.models import resnet18, ResNet18_Weights
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import numpy as np

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')
if device.type == 'cuda':
    print(f'GPU: {torch.cuda.get_device_name(0)}')
print(f'PyTorch version: {torch.__version__}')

Your notebook is powered up and ready to go. Next, you'll load real images from the CIFAR-10 dataset and train your first neural network.

Load CIFAR-10 and Train Your First CNN

Your Colab notebook is set up with a GPU and all the libraries you need. Now it is time to put them to work on real images.

You will download the CIFAR-10 dataset, build a simple convolutional neural network from scratch, and train it. The accuracy you get here becomes the baseline every future improvement gets measured against.

In this step, get ready to:
  • Download the CIFAR-10 dataset and display sample images.
  • Define a 3-layer CNN architecture.
  • Train the model for 15 epochs and track accuracy.
Download CIFAR-10 and visualize sample images

CIFAR-10 is a dataset of 60,000 tiny 32x32 color images spread across 10 categories like airplanes, cats, and trucks. It is a standard benchmark for image classification.

  • Add a new code cell below your setup cell and paste this code to download the dataset and create data loaders:
# ============================================================
# STEP 2: Load CIFAR-10 and Train Simple CNN
# ============================================================

# Basic transforms
transform_basic = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

train_dataset = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=transform_basic
)
test_dataset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=transform_basic
)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False, num_workers=2)

classes = train_dataset.classes
print(f'Classes: {classes}')
print(f'Training samples: {len(train_dataset)}, Test samples: {len(test_dataset)}')

What does this code do?

  • transform_basic converts raw images to tensors and normalizes pixel values from 0-255 to the -1 to 1 range. This helps the model train faster.
  • torchvision.datasets.CIFAR10 downloads the dataset into a ./data folder. The train=True split has 50,000 images and train=False has 10,000 test images.
  • DataLoader wraps each dataset so PyTorch can feed images to the model in batches of 64.
  • Run the cell.
  • You should see the 10 class names printed (airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck) and the counts: 50,000 training samples and 10,000 test samples.

Seeing a download error?

CIFAR-10 downloads from the internet. If the download fails, re-run the cell. Colab occasionally has transient network issues that resolve on retry.

ask me for help debugging a CIFAR-10 download error in Google Colab

Numbers on a screen confirm the data is loaded. Now see what these images actually look like.

  • Add a new code cell and paste this code to display sample images:
# Show sample images
def show_samples(loader, classes):
    images, labels = next(iter(loader))
    fig, axes = plt.subplots(1, 8, figsize=(16, 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]])
        axes[i].axis('off')
    plt.tight_layout()
    plt.show()

show_samples(train_loader, classes)

What does this code do?

show_samples grabs the first batch from the data loader, reverses the normalization (multiplying by 0.5 and adding 0.5 to get pixels back to the 0-1 range), and displays 8 images with their class labels above each one.

  • Run the cell.
  • You should see a row of 8 CIFAR-10 images with class labels like frog, truck, or deer displayed above each image. The images are small and pixelated at 32x32 pixels.

Not seeing any images?

Make sure you ran the data loading cell first. The show_samples function needs train_loader and classes from the previous cell.

help me figure out why my CIFAR-10 sample images are not displaying in Colab

Define a simple 3-layer CNN architecture

Now that you can see the images, it is time to build a model that learns to classify them. You will start with a simple CNN that stacks three convolutional layers.

What is a CNN?

A Convolutional Neural Network (CNN) scans small regions of an image using learnable filters. Each filter detects a specific pattern like an edge, a curve, or a texture.

Stacking multiple convolutional layers lets the network combine simple patterns into complex ones. Early layers detect edges. Later layers detect shapes like wheels or wings built from those edges.

  • Add a new code cell and paste this code to define the CNN class:
# Simple CNN model
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # Three convolutional blocks: Conv2d -> ReLU -> MaxPool2d
        self.conv_layers = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
        )
        # Flatten and classify into 10 categories
        self.fc_layers = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 4 * 4, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )

    def forward(self, x):
        x = self.conv_layers(x)
        x = self.fc_layers(x)
        return x

How does this architecture work?

  • Each convolutional block applies a Conv2d filter, a ReLU activation (keeps only positive values), and a MaxPool2d that halves the image size. After three blocks, a 32x32 image shrinks to 4x4.
  • The fully connected layers (fc_layers) flatten the 4x4 feature maps into a single list of numbers, compress it to 128 values, and output 10 scores. One score per CIFAR-10 class.
  • Run the cell to register the class definition.
Train the model and print epoch-by-epoch results

The CNN architecture is defined, but it has not learned anything yet. Training means feeding images through the model, measuring how wrong it is, and adjusting its filters to be less wrong next time.

You need two helper functions first. One measures accuracy on the test set. The other runs the full training loop.

  • Add a new code cell and paste this evaluation function:
# Evaluate model accuracy on a dataset
def evaluate(model, loader, device):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
    return 100.0 * correct / total

What does evaluate do?

model.eval() switches the model out of training mode. torch.no_grad() tells PyTorch not to track gradients, which saves memory. The function runs every test image through the model, counts how many predictions match the true labels, and returns the percentage.

  • Run the cell to register the function.
  • Add a new code cell and paste this training function:
# Training function - runs the training loop and tracks metrics
def train_model(model, train_loader, test_loader, criterion, optimizer, num_epochs, device):
    train_losses = []
    train_accs = []
    test_accs = []
    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
        train_loss = running_loss / len(train_loader)
        train_acc = 100.0 * correct / total
        test_acc = evaluate(model, test_loader, device)
        train_losses.append(train_loss)
        train_accs.append(train_acc)
        test_accs.append(test_acc)
        print(f'Epoch [{epoch+1}/{num_epochs}] Loss: {train_loss:.4f} '
              f'Train Acc: {train_acc:.2f}% Test Acc: {test_acc:.2f}%')
    return train_losses, train_accs, test_accs

What does train_model do?

  • For each epoch, it loops through every batch of training images. Each batch goes through a forward pass (make predictions), loss calculation (how wrong were they), and a backward pass (adjust the filters).
  • After each epoch, it calls evaluate to measure accuracy on test images the model has never seen. This is the number that actually matters.
  • It stores training loss, training accuracy, and test accuracy in lists so you can plot them later.
  • Run the cell to register the function.

Before you run training, make a prediction. What accuracy do you think a simple 3-layer CNN will achieve when classifying 10 categories of tiny 32x32 images? Random guessing would give you 10%.

  • Add a new code cell and paste this code to create the model, set up the optimizer, and start training:
# Train the simple CNN
simple_model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(simple_model.parameters(), lr=0.001)

print('\nTraining Simple CNN...')
simple_losses, simple_train_accs, simple_test_accs = train_model(
    simple_model, train_loader, test_loader,
    criterion, optimizer, num_epochs=15, device=device
)

What does this code do?

  • SimpleCNN().to(device) creates the model and moves it to the GPU so training runs fast.
  • CrossEntropyLoss measures how wrong the model's predictions are. The lower this number, the better.
  • Adam is the optimizer that adjusts the model's filters after each batch. The learning rate lr=0.001 controls how big each adjustment is.

Training takes about 3-5 minutes on a T4 GPU. You will see output after each of the 15 epochs showing loss, training accuracy, and test accuracy.

When it finishes, the final test accuracy should land around 60-68%. That is much better than the 10% random guessing baseline, but it also means the model gets roughly 1 in 3 images wrong.

Training is extremely slow or crashing?

Check that your GPU is active. In your Colab menu, click Runtime then Change runtime type and confirm T4 GPU is selected. If it says None, your code is running on CPU, which is much slower.

If you see an out-of-memory error, click Runtime then Restart runtime and run all cells from the top.

help me fix slow training or GPU errors in my Colab notebook

✔️ Awesome, I've got everything!

Great. Double check you have saved your notebook. Colab auto-saves, but you can press Ctrl+S (Windows) or Cmd+S (macOS) to be safe.

ⓧ I'd like to double check the full code

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

  • Cell 1: Imports and GPU check (from Step 1).
  • Cell 2: CIFAR-10 download, transforms, data loaders, and class/count prints.
  • Cell 3: show_samples function definition and call.
  • Cell 4: SimpleCNN class definition.
  • Cell 5: evaluate function definition.
  • Cell 6: train_model function definition.
  • Cell 7: Model creation, optimizer setup, and train_model call (15 epochs of output).

Your first CNN is trained, and you have a baseline accuracy to beat. But take a closer look at the training output. Notice how training accuracy climbs much higher than test accuracy? In the next step, you will plot those curves to diagnose exactly what is going wrong.

Diagnose Overfitting and Fight Back

Your simple CNN is trained, but something is off. Look at the numbers from the last step: training accuracy climbed past 85%, yet test accuracy stalled around 65%. That 20-point gap means the model memorized the training images instead of learning general visual patterns.

This is called overfitting, and it is the single biggest obstacle in training neural networks. In this step, you will plot the gap to see it clearly, then fight back with two techniques: data augmentation and dropout.

In this step, get ready to:
  • Plot training curves that expose the overfitting gap in your simple CNN.
  • Build an improved CNN with data augmentation and dropout.
  • Train the improved model and compare its accuracy against the simple CNN.
Plot training curves to expose the overfitting gap

Numbers scrolling by in an epoch log are easy to miss. A plot of training accuracy vs test accuracy makes the overfitting gap impossible to ignore.

Before you run the next cell, take a guess: what do you think the training accuracy curve and the test accuracy curve will look like relative to each other? Will they track together, or diverge?

  • Create a new code cell in your Colab notebook and paste the following code to define a plotting function and visualize the simple CNN's training curves:
# Plot training curves
def plot_training_curves(train_accs, test_accs, title):
    plt.figure(figsize=(10, 5))
    plt.plot(train_accs, label='Train Accuracy')
    plt.plot(test_accs, label='Test Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy (%)')
    plt.title(title)
    plt.legend()
    plt.grid(True)
    plt.show()

plot_training_curves(simple_train_accs, simple_test_accs,
                     'Simple CNN: Train vs Test Accuracy')
print(f'Overfitting gap: {simple_train_accs[-1] - simple_test_accs[-1]:.2f}%')
print('The model is overfitting! Training accuracy is much higher than test accuracy.')

What does this code do?

The plot_training_curves function is a reusable helper that plots two accuracy lines on the same chart. You will call it again later for the improved model.

  • train_accs shows how well the model performs on images it has already seen.
  • test_accs shows how well the model performs on images it has never seen.
  • The gap between these two lines is the overfitting gap. A wider gap means the model is memorizing rather than generalizing.
  • Run the cell.

You should see a chart with two lines. The training accuracy curve climbs steadily toward 85%+, but the test accuracy curve flattens around 65%. That growing gap between the lines is your overfitting problem, visualized.

Below the chart, you should see the exact overfitting gap printed as a percentage (around 20%).

Why is this gap a problem?

High training accuracy with low test accuracy means the model learned the specific pixels of each training image rather than the general shapes that define "airplane" or "cat". When it sees a new image, those memorized pixel patterns do not transfer.

The fix is to make the training data harder to memorize (data augmentation) and to force the model to build redundant internal representations (dropout).

Not seeing the plot?

  • Make sure you ran all previous cells first. The variables simple_train_accs and simple_test_accs are created during the Step 2 training cell.
  • If you see a NameError, your Colab runtime may have restarted. Go to Runtime > Run all to re-execute every cell from the top.

Help me debug why my training curve plot is not appearing in Colab.

Build an improved CNN with data augmentation and dropout

You have seen the overfitting gap. Now you will attack it from two angles. Data augmentation randomly flips and crops training images so the model never sees the exact same image twice. Dropout randomly disables neurons during training, forcing the network to spread its learning across more connections.

  • Create a new code cell and paste the following code to define augmented transforms and reload the training dataset:
# Augmented 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))
])

train_dataset_aug = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=transform_augmented
)
train_loader_aug = DataLoader(train_dataset_aug, batch_size=64, shuffle=True, num_workers=2)

What do these augmentations do?

  • RandomHorizontalFlip() flips each image left-to-right with a 50% chance. A flipped cat is still a cat, but the pixel pattern is different.
  • RandomCrop(32, padding=4) pads each edge by 4 pixels and then crops back to 32x32 at a random position. This shifts the subject around the frame, preventing the model from memorizing positions.
  • The dataset itself does not change. These transforms are applied on-the-fly each time an image is loaded, so every epoch sees slightly different versions of the same images.
  • Run the cell. You will not see any output, but the augmented data pipeline is ready for training.
  • Create a new code cell and paste the following ImprovedCNN class, which adds batch normalization and dropout to the original architecture:
# Improved CNN with BatchNorm and Dropout
class ImprovedCNN(nn.Module):
    def __init__(self):
        super(ImprovedCNN, self).__init__()
        self.conv_layers = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
        )
        self.fc_layers = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 4 * 4, 128),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(128, 10)
        )

    def forward(self, x):
        x = self.conv_layers(x)
        x = self.fc_layers(x)
        return x

What changed from SimpleCNN?

Compare this to your SimpleCNN class from Step 2. The layer sizes and structure are identical, but two new ingredients are added:

  • BatchNorm2d after each Conv2d layer normalizes the values flowing through the network, which stabilizes training and helps the model converge faster.
  • Dropout(0.5) before the final Linear layer randomly shuts off 50% of the neurons during each training pass. This prevents any single neuron from becoming a "crutch" the model relies on.
  • Run the cell. No output is expected. The class is now defined and ready to be instantiated.

Seeing a syntax error?

  • Check that every opening parenthesis has a matching close. The nn.Sequential() blocks have many nested layers, so it is easy to miss a closing parenthesis or comma.
  • Make sure the indentation is consistent. Each layer inside nn.Sequential should be indented one level deeper than the nn.Sequential( line itself.

Help me fix a syntax error in my ImprovedCNN class definition.

Train the improved model and compare results

Your improved CNN is defined and the augmented data pipeline is ready. Before you run the training, make a prediction: do you think augmentation and dropout will fully close the overfitting gap, or just narrow it?

  • Create a new code cell and paste the following code to instantiate the improved model, train it, and compare results:
improved_model = ImprovedCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(improved_model.parameters(), lr=0.001)

print('\nTraining Improved CNN with augmentation and dropout...')
improved_losses, improved_train_accs, improved_test_accs = train_model(
    improved_model, train_loader_aug, test_loader,
    criterion, optimizer, num_epochs=15, device=device
)

plot_training_curves(improved_train_accs, improved_test_accs,
                     'Improved CNN: Train vs Test Accuracy')

print(f'\nComparison:')
print(f'  Simple CNN test accuracy:   {simple_test_accs[-1]:.2f}%')
print(f'  Improved CNN test accuracy: {improved_test_accs[-1]:.2f}%')

What does this code do?

  • The improved model is instantiated and moved to the GPU, just like the simple model was in Step 2.
  • It trains for 15 epochs using the same train_model function, but this time it receives train_loader_aug (the augmented data) instead of the original train_loader.
  • After training, it plots the new training curves and prints a side-by-side comparison of both models.
  • Run the cell. This will take a few minutes as the model trains for 15 epochs.

You should see epoch-by-epoch results printing, followed by a training curve plot for the improved CNN. The gap between the train and test curves is narrower than before.

The final comparison should show your simple CNN at around 65% test accuracy and the improved CNN at around 72-76%. An improvement, but still not great.

Why didn't augmentation and dropout get us to 90%?

Augmentation and dropout reduced overfitting, but they did not give the model better eyes. A 3-layer CNN trained from scratch on 50,000 tiny images simply does not learn enough visual patterns to distinguish 10 classes well.

To break past 75%, you need a model that already knows what edges, textures, and shapes look like. That is exactly what transfer learning provides in the next step.

Training seems stuck or accuracy is very low?

  • If accuracy stays below 30%, check that you are passing train_loader_aug (not train_loader) to the train_model function.
  • If you see a RuntimeError about tensor sizes, make sure the ImprovedCNN class cell ran successfully before this cell.
  • If the Colab runtime disconnected during training, reconnect and use Runtime > Run all to re-execute everything from the start.

Help me figure out why my improved CNN is not training correctly.

✔️ Awesome, I've got everything!

Great. Make sure all four cells from this step ran without errors before moving on.

ⓧ I'd like to double check the full code

Your notebook should now have the following cells in order from top to bottom. Scroll through and confirm each one is present:

  • Step 1 cell: Imports, device setup, and PyTorch version print.
  • Step 2 cells: transform_basic, dataset loading, show_samples, SimpleCNN class, train_model and evaluate functions, and simple model training.
  • Step 3 cell 1: plot_training_curves function, simple CNN plot call, and overfitting gap print.
  • Step 3 cell 2: transform_augmented, train_dataset_aug, and train_loader_aug.
  • Step 3 cell 3: ImprovedCNN class definition with BatchNorm2d and Dropout(0.5).
  • Step 3 cell 4: Improved model instantiation, training, curve plot, and comparison print.

You went from ~65% to ~73% by making training harder and forcing redundancy. That is real progress. But a CNN trained from scratch on 50,000 small images has a ceiling. In the next step, you will blow past it by importing millions of visual patterns that a much larger model already learned.

Supercharge with Transfer Learning

Your improved CNN narrowed the overfitting gap and pushed test accuracy to around 73%. That is real progress from the 65% baseline.

But 73% still means the model gets roughly 1 in 4 images wrong. A CNN trained from scratch on 50,000 small images simply does not learn enough visual patterns to do better. Transfer learning solves this by starting from a model that has already learned millions of patterns from a much larger dataset.

In this step, get ready to:
  • Prepare CIFAR-10 data for ResNet18's input requirements.
  • Load a pretrained ResNet18 model and replace its final layer for 10 classes.
  • Train the new classification layer and compare all three models.
Prepare the data for ResNet18

ResNet18 was trained on 224x224 ImageNet images. CIFAR-10 images are only 32x32 pixels. You need to resize them and normalize with the exact same statistics ImageNet used, so ResNet18's pretrained patterns interpret the pixel values correctly.

  • Create a new code cell with the ResNet data transforms by adding this code:
# ResNet transforms with ImageNet normalization
transform_resnet_train = transforms.Compose([
    transforms.Resize(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

transform_resnet_test = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

What does this code do?

  • Resize(224) scales each 32x32 CIFAR-10 image up to 224x224, the size ResNet18 expects.
  • Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) applies the exact normalization values ResNet18 was trained with on ImageNet. Using different values would confuse the pretrained layers.
  • The training transforms include RandomHorizontalFlip() for data augmentation. The test transforms skip augmentation so evaluation stays consistent.

Why different normalization values?

Your earlier CNNs used (0.5, 0.5, 0.5) for both mean and std. That was a simple rescaling to the [-1, 1] range. ResNet18's pretrained weights were trained with ImageNet's actual channel-by-channel statistics, so you must match them exactly.

  • Run the cell. It defines the transform pipelines for the next cell to use.
  • Load CIFAR-10 with the ResNet transforms and create DataLoaders by creating a new code cell with this code:
train_dataset_resnet = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=transform_resnet_train
)
test_dataset_resnet = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=transform_resnet_test
)

train_loader_resnet = DataLoader(
    train_dataset_resnet, batch_size=32, shuffle=True, num_workers=2
)
test_loader_resnet = DataLoader(
    test_dataset_resnet, batch_size=32, shuffle=False, num_workers=2
)

What does this code do?

  • This reloads the same CIFAR-10 data you used before, but with the new ResNet transforms applied. The cached data is reused, so no fresh download happens.
  • batch_size=32 is smaller than the 64 you used earlier. Each image is now 224x224 instead of 32x32, so each batch uses more GPU memory.
  • Run the cell. You should see Files already downloaded and verified printed for both datasets.

Seeing a memory error?

If you see a CUDA out of memory error, your GPU may be holding memory from earlier training. Go to Runtime > Restart runtime, then re-run all your cells from the top.

Help me fix a CUDA out of memory error in my Colab notebook.

Load pretrained ResNet18 and replace the final layer

ResNet18 was trained on 1.2 million ImageNet images across 1,000 categories. Its convolutional layers already know how to detect edges, textures, shapes, and object parts. Instead of learning all of that from scratch, you will keep those pretrained layers frozen and only replace the final classification layer to output 10 CIFAR-10 classes.

  • Load and configure the pretrained model by creating a new code cell with this code:
# Load pretrained ResNet18
resnet_model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)

# Freeze all layers
for param in resnet_model.parameters():
    param.requires_grad = False

# Replace the final fully connected layer for 10 classes
num_features = resnet_model.fc.in_features
resnet_model.fc = nn.Linear(num_features, 10)
resnet_model = resnet_model.to(device)

What does this code do?

  • resnet18(weights=ResNet18_Weights.IMAGENET1K_V1) downloads the model with pretrained weights from ImageNet. These weights encode everything the model learned about visual patterns.
  • param.requires_grad = False freezes every layer so the pretrained patterns are not changed during training.
  • resnet_model.fc = nn.Linear(num_features, 10) replaces the original 1,000-class output layer with a fresh 10-class layer for CIFAR-10. This new layer's requires_grad is True by default, so it is the only part that will learn.
  • Run the cell. You should see a download progress bar as PyTorch fetches the pretrained weights (about 45 MB).

Why freeze layers instead of retraining everything?

This approach is called feature extraction. ResNet18's convolutional layers already know how to extract useful features from images. Training only the final classification layer is faster (5 epochs instead of 15+) and avoids accidentally overwriting the pretrained patterns with noise from a relatively small dataset.

Model not downloading?

Make sure your Colab runtime is connected. Check the top-right corner for a green checkmark. If you see a disconnected icon, click Reconnect and re-run the cell.

Help me fix a model download issue in Google Colab.

Train only the new classification layer and compare all three models

Everything is set. The pretrained layers are locked, and the new classification head is ready to learn. You will train for just 5 epochs because the model already starts with strong features.

Before you run this, what accuracy do you think ResNet18 will hit compared to the improved CNN's ~73%? Will the jump be small or dramatic?

  • Train the ResNet18 model and print the final comparison by creating a new code cell with this code:
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(resnet_model.fc.parameters(), lr=0.001)

print('\nTraining ResNet18 with Transfer Learning...')
resnet_losses, resnet_train_accs, resnet_test_accs = train_model(
    resnet_model, train_loader_resnet, test_loader_resnet,
    criterion, optimizer, num_epochs=5, device=device
)

print(f'\n--- Final Results ---')
print(f'  Simple CNN:         {simple_test_accs[-1]:.2f}%')
print(f'  Improved CNN:       {improved_test_accs[-1]:.2f}%')
print(f'  ResNet18 Transfer:  {resnet_test_accs[-1]:.2f}%')

What does this code do?

  • resnet_model.fc.parameters() tells the optimizer to update only the new classification layer. The frozen pretrained layers are untouched.
  • num_epochs=5 is enough because the model already understands image features. It only needs to learn how to map those features to 10 CIFAR-10 classes.
  • The final print block shows all three model accuracies side by side so you can compare the progression.

Training takes around 10-15 minutes

Each image is now 224x224 instead of 32x32, so each batch takes longer to process. You will see epoch-by-epoch progress printing as it trains. This is a good time to review your training curve plots from Step 3.

  • Run the cell and watch the epoch-by-epoch output.

Your final results should show ResNet18 hitting 90%+ test accuracy. Compare that to the Simple CNN at ~65% and the Improved CNN at ~73%. That is the power of transfer learning. Five epochs and one new layer outperform everything you built from scratch.

Accuracy below 85%?

Make sure you froze the pretrained layers with param.requires_grad = False before replacing the final layer. If you froze them after, the new fc layer got frozen too and cannot learn.

Also check that your optimizer targets resnet_model.fc.parameters() and not resnet_model.parameters(). Passing all parameters when they are frozen means nothing trains at all.

Help me debug low transfer learning accuracy in my ResNet18 model.

✔️ Awesome, I've got everything!

Great. Make sure all four code cells from this step are saved in your notebook.

ⓧ I'd like to double check the full code

Your notebook should have the following four cells from this step, in order. Scroll up and confirm each one is present:

  • ResNet transforms cell: defines transform_resnet_train and transform_resnet_test with Resize(224) and ImageNet normalization.
  • Datasets and DataLoaders cell: creates train_dataset_resnet, test_dataset_resnet, train_loader_resnet, and test_loader_resnet with batch_size=32.
  • Model setup cell: loads resnet_model with ResNet18_Weights.IMAGENET1K_V1, freezes all parameters, replaces resnet_model.fc with nn.Linear(num_features, 10).
  • Training and comparison cell: trains for 5 epochs with resnet_model.fc.parameters() and prints the final three-model comparison.

Your ResNet18 model just jumped past 90% accuracy in only 5 epochs. Next up, you will evaluate exactly where it succeeds and fails using a confusion matrix and per-class accuracy breakdown.

Evaluate Your Model Like a Pro

Your ResNet18 model just hit 90%+ accuracy on CIFAR-10. That is a massive leap from the 65% baseline. But a single accuracy percentage hides a lot of detail.

Which classes does the model still confuse? Are cats harder to classify than trucks? In this step, you will build three evaluation tools that answer these questions: a confusion matrix heatmap, a color-coded prediction grid, and a per-class accuracy breakdown.

In this step, get ready to:
  • Generate a confusion matrix heatmap to see which classes the model mixes up.
  • Display sample predictions with green (correct) and red (incorrect) color-coded labels.
  • Print a per-class accuracy breakdown for all 10 CIFAR-10 classes.
Generate a confusion matrix heatmap

A confusion matrix is a grid that shows how often the model predicted each class correctly and how often it confused one class for another. Each row represents the actual class, and each column represents what the model predicted.

Before you can build the matrix, you need to run the trained resnet_model on every test image and collect its predictions.

  • Create a new code cell and import the evaluation libraries, then collect all predictions by pasting this code:
from sklearn.metrics import confusion_matrix
import seaborn as sns

# Collect all predictions
resnet_model.eval()
all_preds = []
all_labels = []
with torch.no_grad():
    for images, labels in test_loader_resnet:
        images = images.to(device)
        outputs = resnet_model(images)
        _, predicted = outputs.max(1)
        all_preds.extend(predicted.cpu().numpy())
        all_labels.extend(labels.numpy())

What does this code do?

  • confusion_matrix from scikit-learn computes how many times each class was predicted as each other class.
  • seaborn is a visualization library that makes clean, colorful heatmaps from data.
  • The loop runs every test image through the model with gradients disabled (torch.no_grad()) and stores the predicted class and true label for all 10,000 images.
  • Run the cell. You will not see any output yet. The next chunk uses these predictions to build the heatmap.

Now that the predictions are collected, think about this: which classes do you think the model confuses most often? Animals that look alike, or completely different categories?

  • Create a new code cell and compute the confusion matrix heatmap by pasting this code:
# Confusion matrix heatmap
cm = confusion_matrix(all_labels, all_preds)
plt.figure(figsize=(10, 8))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
            xticklabels=classes, yticklabels=classes)
plt.xlabel('Predicted')
plt.ylabel('Actual')
plt.title('Confusion Matrix - ResNet18 Transfer Learning')
plt.tight_layout()
plt.show()

What does this code do?

  • confusion_matrix(all_labels, all_preds) builds a 10x10 grid where each cell counts how many images of the actual class (row) the model predicted as a specific class (column).
  • sns.heatmap turns that grid into a color-coded chart. Darker blue squares mean more images fell into that cell. The diagonal shows correct predictions.
  • annot=True prints the exact count inside each cell, and fmt='d' formats them as whole numbers.
  • Run the cell.

You should see a 10x10 heatmap with class names on both axes. The diagonal cells are darkest blue because those are correct predictions. Off-diagonal cells reveal where the model confuses classes.

What should I look for?

Look at the off-diagonal cells with the highest numbers. You will likely see the model confusing cat and dog more than any other pair. That makes sense. Both are four-legged furry animals in small 32x32 images.

Not seeing the heatmap?

  • Make sure you ran the previous cell first. The all_preds and all_labels lists must be populated before computing the matrix.
  • Check that the imports for confusion_matrix and seaborn are in the previous cell, not missing.
  • Why isn't my confusion matrix heatmap showing up in Colab?
Display sample predictions with correct and incorrect indicators

The confusion matrix tells you which classes get mixed up. Now you need to see the actual images behind those numbers. A prediction grid shows real test images with color-coded labels so you can spot patterns in what the model gets right and wrong.

  • Create a new code cell and define the prediction display function by pasting this code:
# Sample predictions with color-coded labels
def show_predictions(model, dataset, classes, device, num=10):
    model.eval()
    fig, axes = plt.subplots(2, 5, figsize=(15, 6))
    axes = axes.flatten()
    for i in range(num):
        image, label = dataset[i]
        with torch.no_grad():
            output = model(image.unsqueeze(0).to(device))
            _, predicted = output.max(1)

        img = image.clone()
        mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
        std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
        img = img * std + mean
        img = torch.clamp(img, 0, 1)

        color = 'green' if predicted.item() == label else 'red'
        axes[i].imshow(img.permute(1, 2, 0).numpy())
        axes[i].set_title(
            f'Pred: {classes[predicted.item()]}\nTrue: {classes[label]}',
            color=color, fontsize=9
        )
        axes[i].axis('off')
    plt.suptitle('Sample Predictions (Green=Correct, Red=Wrong)', fontsize=13)
    plt.tight_layout()
    plt.show()

show_predictions(resnet_model, test_dataset_resnet, classes, device)

What does this code do?

  • show_predictions picks 10 test images and runs each one through the model to get a predicted class.
  • The image needs to be "un-normalized" before displaying. The model received images scaled by ImageNet mean and standard deviation, so the function reverses that (img * std + mean) to get pixel values back to the 0-1 range.
  • Each image title shows both the predicted and true class. Green text means the model got it right. Red means it got it wrong.
  • Run the cell.

You should see a 2x5 grid of 10 test images. Each image has a title showing the predicted label and the true label. Correct predictions appear in green text. Incorrect ones appear in red.

Seeing all-black or distorted images?

  • Make sure the mean and std values in the function match the ImageNet normalization values used in your transforms: mean=[0.485, 0.456, 0.406] and std=[0.229, 0.224, 0.225].
  • Check that torch.clamp(img, 0, 1) is present. Without it, pixel values outside 0-1 cause display artifacts.
  • Why are my prediction images displaying incorrectly in my Colab notebook?
Print per-class accuracy breakdown

The confusion matrix and prediction grid give you a visual read. Now you need the exact numbers. A per-class accuracy breakdown tells you the precise percentage for each of the 10 categories.

Before you run this, predict: which class do you think the model handles worst? Think about what you saw in the confusion matrix.

  • Create a new code cell and compute per-class accuracy by pasting this code:
# Per-class accuracy
class_correct = [0] * 10
class_total = [0] * 10
for pred, label in zip(all_preds, all_labels):
    class_correct[label] += int(pred == label)
    class_total[label] += 1

print('\nPer-class accuracy:')
for i in range(10):
    acc = 100.0 * class_correct[i] / class_total[i]
    print(f'  {classes[i]:12s}: {acc:.1f}%')

What does this code do?

  • class_correct and class_total are lists of 10 counters, one per class. The loop walks through every prediction and tallies correct guesses per class.
  • The final print loop divides correct by total for each class and formats the result as a tidy percentage table. This tells you exactly where the model excels and where it still struggles.
  • Run the cell.

You should see a formatted table listing each of the 10 CIFAR-10 classes with its accuracy percentage. Categories like ship and automobile typically score highest. Categories like cat tend to score lowest because they share visual features with dog.

Seeing a NameError?

  • If you see an error about all_preds or all_labels not being defined, scroll up and re-run the first cell in this step (the one that imports confusion_matrix and collects predictions).
  • If you restarted the Colab runtime at any point, you need to re-run all previous cells from Step 1 onward.
  • Why am I getting a NameError for all_preds or all_labels in my Colab notebook?

✔️ Awesome, I've got everything!

Great. You should now have four new cells in your notebook for this step: the imports and prediction collection, the confusion matrix heatmap, the sample prediction grid, and the per-class accuracy table.

ⓧ I'd like to double check the full code

Scroll through the cells you added in this step and confirm each one is present, in order:

  • Cell 1: Imports (from sklearn.metrics import confusion_matrix and import seaborn as sns) plus the prediction collection loop that populates all_preds and all_labels.
  • Cell 2: Confusion matrix computation (cm = confusion_matrix(all_labels, all_preds)) and sns.heatmap display with Blues colormap.
  • Cell 3: show_predictions function definition and the call show_predictions(resnet_model, test_dataset_resnet, classes, device).
  • Cell 4: Per-class accuracy loop using class_correct and class_total lists, printing accuracy for all 10 classes.

Secret mission

See Through Your Model's Eyes with Grad-CAM

Your model hits 90%+ accuracy, but can you explain why it makes each prediction? In this bonus challenge, you'll install the Grad-CAM library and generate heatmap visualizations that reveal exactly which pixels your ResNet18 model focuses on when classifying each image.

Clean Up Your Resources

Clean Up Your Resources

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

Resources you used:

  • Google Colab notebook (saved in your Google Drive).
  • CIFAR-10 dataset cache (stored in Colab's temporary runtime storage at ./data).
  • Trained model weights (held in Colab's runtime memory).

Keep everything running

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

  • Your notebook stays in Google Drive and is accessible anytime from colab.research.google.com.
  • Colab automatically disconnects idle runtimes after a period of inactivity, so the GPU will free itself.
  • The CIFAR-10 dataset and trained model weights live in the temporary runtime. They will be cleared when the runtime disconnects, but your notebook code is preserved in Google Drive.

Pause - I'll come back to this later

Disconnect the GPU runtime to free up resources for other Colab users. Your notebook code stays safe in Google Drive.

  • In your Colab notebook, click Runtime in the top menu bar.
  • Click Disconnect and delete runtime.
  • This clears the downloaded CIFAR-10 data and trained models from memory. When you reconnect later, you will need to re-run the cells to retrain.

Delete - I don't want to use this again

Remove your notebook and disconnect the runtime to fully clean up.

  • Disconnect the runtime first. Click Runtime in the top menu bar, then click Disconnect and delete runtime. This removes the CIFAR-10 dataset cache and all trained model weights from Colab's temporary storage.
  • Go to drive.google.com and find your notebook in Colab Notebooks.
  • Right-click the notebook file and select Move to trash.
  • Empty your Google Drive trash to permanently delete the file. Click Trash in the left sidebar, then click Empty trash.

Nice Work!

Nice Work!

You just built a complete image classification pipeline in PyTorch that classifies real-world images across 10 categories. You went from a basic CNN struggling at ~65% accuracy all the way to a transfer learning model hitting 90%+ on CIFAR-10.

You've learned how to:

  • Train three progressively better CNN models on CIFAR-10, from a simple 3-layer baseline at ~65% accuracy, to an improved CNN with data augmentation and dropout at ~73%, to a pretrained ResNet18 achieving 90%+.
  • Diagnose overfitting by plotting training curves that reveal the gap between train and test accuracy, then fight it with batch normalization, dropout, and augmented training data.
  • Evaluate model performance using a confusion matrix heatmap, a per-class accuracy breakdown, and a color-coded prediction grid that shows exactly where your model succeeds and fails.
  • Secret Mission: Generated Grad-CAM heatmap visualizations to see exactly which image regions the model focuses on when making predictions.

Ready to quiz yourself?