Build an Image Classifier with PyTorch

Train a CNN and fine-tune ResNet18 to classify images with 90%+ accuracy.

Introduction

30 Second Summary

You can tell the difference between a cat and a dog in a split second. Teaching a computer to do the same thing requires thousands of example images and a structured training process.

In this project, you will build an image classifier that learns to recognize 10 categories of objects. You will train two neural networks using PyTorch on a free cloud GPU, then compare their performance to see why transfer learning is so powerful.

What You'll Build

A complete image classification pipeline in Google Colab that trains two models on the CIFAR-10 dataset (60,000 images across 10 categories like airplane, car, bird, and cat).

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

  • A GPU-accelerated deep learning environment running in Google Colab with PyTorch.
  • A custom CNN built from scratch, trained to ~65-70% accuracy on 10 image categories.
  • A pretrained ResNet18 fine-tuned to achieve 90%+ accuracy, proving the power of transfer learning.
  • Secret Mission: Upload and classify your own images with the trained model.

Are there any prerequisites?

You should be comfortable with basic Python (variables, loops, functions). Familiarity with NumPy is helpful but not required.

Set Up Google Colab with GPU

Before you can train a neural network, you need a computing environment with serious processing power. Training on a regular CPU can take hours, but a GPU crunches through the same work in minutes.

Google Colab gives you free access to a GPU right in your browser. In this step, you will create a Colab notebook, enable GPU acceleration, and confirm that PyTorch can see the GPU.

In this step, get ready to:
  • Create a new Google Colab notebook.
  • Enable T4 GPU acceleration.
  • Verify that PyTorch detects the GPU.
Open Google Colab and create a new notebook
  • Go to colab.research.google.com in your browser.
  • Click New Notebook in the bottom left of the pop-up window.
  • Click the notebook title at the top left (it defaults to something like Untitled0.ipynb).
  • Rename it to CIFAR10_Image_Classifier.

You should see your new notebook open with an empty code cell, ready for input.

What is Google Colab?

Google Colab is a free, cloud-based Jupyter notebook environment. It comes with popular Python libraries like PyTorch pre-installed, so you can start coding immediately without any local setup.

Enable GPU acceleration

By default, Colab runs your code on a CPU. You need to switch the runtime to use a GPU so training runs quickly.

  • Click Runtime in the top menu bar.
  • Select Change runtime type.
  • Under Hardware accelerator, select T4 GPU.
  • Click Save.

Why T4 GPU?

The T4 is a powerful GPU that Colab provides for free. It has enough memory and processing power to train the models in this project in just a few minutes instead of hours on a CPU.

Verify GPU access and import libraries

Now you will import all the libraries the project needs and confirm that PyTorch can see the GPU.

  • Click the empty code cell in your notebook and paste the following code:
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

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

What does this code do?

The first section imports the libraries you will use throughout the project:

  • torch is the core PyTorch library for building and training neural networks.
  • torchvision provides datasets and image transformation tools.
  • matplotlib and numpy help with visualization and numerical operations.

The device variable checks whether CUDA (NVIDIA's GPU computing platform) is available. If it is, PyTorch will use the GPU for all heavy computations.

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

You should see output showing the PyTorch version and confirming that CUDA is available with a GPU device name.

Seeing 'cpu' instead of 'cuda'?

If the output says Using device: cpu, the GPU is not enabled. Go back to Runtime > Change runtime type and make sure T4 GPU is selected. After changing it, re-run the cell.

If GPU is not available as an option, Colab may be at capacity. Wait a few minutes and try again.

Why isn't my Colab notebook detecting a GPU?

✔️ Awesome, I've got everything!

Great. Double check you see output similar to this:

PyTorch version: 2.10.0+cu124
Using device: cuda
GPU: Tesla T4

Your exact version string and GPU name may differ slightly, but the key thing is that device shows cuda and a GPU name is printed.

ⓧ I'd like to double check the full code

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

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

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

Your GPU-powered environment is ready. Next up, you will load the CIFAR-10 dataset and prepare it for training.

Load and Preprocess CIFAR-10

Your GPU environment is confirmed and ready to go. Now you need training data for your neural network to learn from.

CIFAR-10 is a classic computer vision benchmark containing 60,000 images across 10 categories like airplane, cat, and truck. In this step, you'll download the dataset, preprocess it for PyTorch, and visualize samples to confirm everything loaded correctly.

In this step, get ready to:
  • Define a preprocessing pipeline that normalizes image pixel values.
  • Download the CIFAR-10 training and test sets with DataLoaders.
  • Visualize a grid of sample images to verify the data.
Define transforms and load the dataset

Before feeding images to a neural network, you need to convert them to tensors and normalize their pixel values. Normalization keeps input values in a consistent range, which helps the model train faster and more stably.

  • Create a new code cell below your GPU verification cell.
  • Add the following code to define your transform pipeline, download CIFAR-10, and create data loaders:
# Define transforms
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
])

# Download and load datasets
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64,
                                          shuffle=True, num_workers=2)

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

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

print(f"Training samples: {len(trainset)}")
print(f"Test samples: {len(testset)}")
print(f"Number of classes: {len(classes)}")

What does this code do?

  • transforms.Compose chains multiple preprocessing steps into a single pipeline.
  • ToTensor() converts images from PIL format into PyTorch tensors with pixel values between 0 and 1.
  • Normalize() shifts pixel values using CIFAR-10's per-channel mean and standard deviation. This centers the data around zero so the model can learn more efficiently.
  • DataLoader wraps the dataset to serve images in batches of 64 during training.
  • shuffle=True randomizes the training order each epoch so the model does not memorize sequence patterns.
  • Run the cell by pressing Shift+Enter.
  • You should see download progress for CIFAR-10, followed by output confirming 50,000 training samples, 10,000 test samples, and 10 classes.

Not seeing the download output?

  • Make sure you ran Cell 1 (the imports cell) first. Cell 2 depends on the libraries imported there.
  • If you see a connection error, check that your Colab runtime is still active. Click the RAM/Disk indicator in the top right to verify.

Help me debug my CIFAR-10 data loading cell in Google Colab.

Why these specific normalization values?

The mean (0.4914, 0.4822, 0.4465) and standard deviation (0.2470, 0.2435, 0.2616) are pre-computed statistics from the CIFAR-10 training set, one value per RGB color channel.

Using these dataset-specific values (instead of generic 0.5) gives the model a head start by matching the actual distribution of pixel intensities in CIFAR-10.

Visualize a batch of training images

Visualizing your data confirms the images loaded correctly and helps you understand what the model will be working with. CIFAR-10 images are tiny (32x32 pixels), so they will look a bit blurry.

  • Create a new code cell below the data loading cell.
  • Add the following code to display a grid of sample training images with their labels:
def imshow(img):
    # Unnormalize
    img = img * torch.tensor([0.2470, 0.2435, 0.2616]).view(3, 1, 1) + torch.tensor([0.4914, 0.4822, 0.4465]).view(3, 1, 1)
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.axis('off')

# Get a batch of training images
dataiter = iter(trainloader)
images, labels = next(dataiter)

# Show images
fig, axes = plt.subplots(2, 8, figsize=(16, 4))
for i, ax in enumerate(axes.flat):
    img = images[i]
    img = img * torch.tensor([0.2470, 0.2435, 0.2616]).view(3, 1, 1) + torch.tensor([0.4914, 0.4822, 0.4465]).view(3, 1, 1)
    img = img.clamp(0, 1)
    ax.imshow(np.transpose(img.numpy(), (1, 2, 0)))
    ax.set_title(classes[labels[i]], fontsize=10)
    ax.axis('off')
plt.tight_layout()
plt.show()

What does this code do?

  • The unnormalization reverses the Normalize() transform so images display with their original colors instead of shifted values.
  • iter(trainloader) and next() grab one batch of 64 images from the training set.
  • The 2x8 subplot grid displays 16 of those images, each labeled with its class name.
  • clamp(0, 1) clips any pixel values that drift outside the valid display range after unnormalization.
  • Run the cell by pressing Shift+Enter.
  • You should see a 2x8 grid of small color images (32x32 pixels), each labeled with a class name like "cat", "ship", or "frog".

Images look weird or all one color?

  • If images appear with strange colors, double-check that the unnormalization values in Cell 3 match the normalization values in Cell 2. They should be the same numbers.
  • If you get a NameError for trainloader or classes, make sure Cell 2 ran successfully before running Cell 3.

Help me fix my CIFAR-10 image visualization.

✔️ Awesome, I've got everything!

Great. Make sure both Cell 2 and Cell 3 ran without errors before moving on.

ⓧ I'd like to double check the full code

Here is the full code for Cell 2:

# Define transforms
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
])

# Download and load datasets
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64,
                                          shuffle=True, num_workers=2)

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

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

print(f"Training samples: {len(trainset)}")
print(f"Test samples: {len(testset)}")
print(f"Number of classes: {len(classes)}")

And here is the full code for Cell 3:

def imshow(img):
    # Unnormalize
    img = img * torch.tensor([0.2470, 0.2435, 0.2616]).view(3, 1, 1) + torch.tensor([0.4914, 0.4822, 0.4465]).view(3, 1, 1)
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.axis('off')

# Get a batch of training images
dataiter = iter(trainloader)
images, labels = next(dataiter)

# Show images
fig, axes = plt.subplots(2, 8, figsize=(16, 4))
for i, ax in enumerate(axes.flat):
    img = images[i]
    img = img * torch.tensor([0.2470, 0.2435, 0.2616]).view(3, 1, 1) + torch.tensor([0.4914, 0.4822, 0.4465]).view(3, 1, 1)
    img = img.clamp(0, 1)
    ax.imshow(np.transpose(img.numpy(), (1, 2, 0)))
    ax.set_title(classes[labels[i]], fontsize=10)
    ax.axis('off')
plt.tight_layout()
plt.show()

Your data pipeline is in place and you have confirmed the images look correct. Next up, you'll build a convolutional neural network from scratch and train it to classify these 10 categories.

Build and Train a CNN from Scratch

Your dataset is loaded and you've confirmed the images look correct. Now it's time to build the neural network that will learn to tell airplanes from frogs.

A Convolutional Neural Network (CNN) is a type of neural network designed specifically for image data. Instead of looking at every pixel independently, it scans for visual patterns like edges, textures, and shapes using small sliding filters. In this step, you'll define a CNN architecture and train it on the 50,000 training images.

In this step, get ready to:
  • Define a CNN with convolutional and fully connected layers.
  • Train the model for 10 epochs on the CIFAR-10 training set.
  • Track test accuracy as the model learns.
Define a simple CNN architecture

Your CNN will have two types of layers. Convolutional layers scan the image for patterns using small filters. Fully connected layers take those detected patterns and map them to the 10 output classes.

  • Create a new code cell in your notebook and paste the following code:
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 8 * 8, 512)
        self.fc2 = nn.Linear(512, 128)
        self.fc3 = nn.Linear(128, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.view(-1, 64 * 8 * 8)
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = SimpleCNN().to(device)
print(model)
total_params = sum(p.numel() for p in model.parameters())
print(f"\nTotal parameters: {total_params:,}")

What does this code do?

  • conv1 takes the 3 color channels (RGB) and produces 32 feature maps. conv2 takes those 32 maps and produces 64.
  • pool (MaxPool2d) halves the image dimensions after each convolution, shrinking 32x32 to 16x16 to 8x8.
  • The three fully connected layers (fc1, fc2, fc3) gradually compress the information from 64*8*8 = 4096 values down to 10 class scores.
  • .to(device) moves the model onto the GPU so training runs fast.
  • Run the cell by pressing the play button or hitting Shift+Enter.
  • You should see the full model architecture printed, followed by the total parameter count (around 2 million parameters).

Seeing a CUDA error?

If you see a message about CUDA not being available, your runtime may have disconnected from the GPU. Click Runtime in the top menu, then Change runtime type, and confirm T4 GPU is still selected.

help me fix a CUDA device error in my CNN model

Set up training components

A model by itself doesn't learn anything. It needs two things to train: a way to measure how wrong its predictions are, and a strategy for adjusting its weights to be less wrong next time.

What are loss functions and optimizers?

CrossEntropyLoss is the standard loss function for classification tasks. It measures how far the model's predicted probabilities are from the correct answer.

Adam is an optimizer that adjusts model weights after each batch. The learning rate (0.001) controls how big each adjustment is. Too large and the model overshoots. Too small and it learns too slowly.

An epoch is one full pass through the entire training set. Training for 10 epochs means the model sees all 50,000 images ten times.

Run the training loop

The training loop is the core pattern in every PyTorch project. For each batch of images, the model makes predictions, calculates how wrong it was, and updates its weights to improve.

  • Create a new code cell and start training by pasting this code:
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

num_epochs = 10

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

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

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

        running_loss += loss.item()
        if (i + 1) % 100 == 0:
            print(f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(trainloader)}], Loss: {running_loss/100:.4f}")
            running_loss = 0.0

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

    print(f"Epoch [{epoch+1}/{num_epochs}] Test Accuracy: {100 * correct / total:.2f}%\n")

print("Training complete!")

What does the training loop do?

  • optimizer.zero_grad() resets the gradients from the previous batch so they don't accumulate.
  • loss.backward() calculates how much each weight contributed to the error (backpropagation).
  • optimizer.step() updates all the model's weights based on those gradients.
  • After each epoch, the loop switches to model.eval() and tests accuracy on the 10,000 test images the model has never seen.
  • Run the cell. Training takes approximately 3-4 minutes on a T4 GPU.
  • You should see loss values printing every 100 steps, and test accuracy printed at the end of each epoch.
  • By the final epoch, test accuracy should reach approximately 65-70%.

What should I see?

The loss values should decrease over time, starting around 1.4-1.5 and dropping below 1.0 by the later epochs. Test accuracy starts low (around 40-50% in epoch 1) and gradually climbs to 65-70%. This is your model learning.

Training seems stuck or accuracy stays low?

If accuracy stays below 20% after several epochs, the model may not be on the GPU. Check that your Cell 1 printed a GPU device name. If it shows CPU, reconnect to a GPU runtime.

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

help me debug my CNN training loop

✔️ Awesome, I've got everything!

Great. Make sure both cells ran without errors before moving on.

ⓧ I'd like to double check the full code

Here is Cell 4 (Define the CNN):

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 8 * 8, 512)
        self.fc2 = nn.Linear(512, 128)
        self.fc3 = nn.Linear(128, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.view(-1, 64 * 8 * 8)
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = SimpleCNN().to(device)
print(model)
total_params = sum(p.numel() for p in model.parameters())
print(f"\nTotal parameters: {total_params:,}")

Here is Cell 5 (Train the CNN):

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

num_epochs = 10

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

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

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

        running_loss += loss.item()
        if (i + 1) % 100 == 0:
            print(f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(trainloader)}], Loss: {running_loss/100:.4f}")
            running_loss = 0.0

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

    print(f"Epoch [{epoch+1}/{num_epochs}] Test Accuracy: {100 * correct / total:.2f}%\n")

print("Training complete!")

Your CNN is trained and can classify images with 65-70% accuracy. That's a solid start, but some categories are clearly harder than others. Next up, you'll evaluate exactly where the model succeeds and where it struggles.

Evaluate the From-Scratch Model

Your CNN is trained and scoring around 65-70% accuracy on the test set. That overall number is a useful starting point, but it hides important details.

In this step, you'll dig into per-class metrics to see which categories your model handles well and which ones trip it up. You'll also visualize actual predictions so you can see the model's mistakes with your own eyes.

In this step, get ready to:
  • Compute per-class accuracy to find the model's weak spots.
  • Visualize predictions with color-coded labels showing correct vs incorrect.
Compute per-class and overall accuracy

An overall accuracy of 68% might sound decent, but what if the model gets 90% on ships and only 40% on cats? Per-class accuracy reveals these imbalances so you know exactly where the model is weak.

  • Create a new code cell below your training cell.
  • Paste the following code to compute accuracy for each of the 10 classes:
model.eval()
class_correct = [0] * 10
class_total = [0] * 10

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

print("Per-class accuracy (From-Scratch CNN):")
print("-" * 40)
for i in range(10):
    acc = 100 * class_correct[i] / class_total[i]
    print(f"  {classes[i]:12s}: {acc:.1f}%")

overall_acc_cnn = 100 * sum(class_correct) / sum(class_total)
print(f"\n  Overall: {overall_acc_cnn:.1f}%")

What does this code do?

model.eval() switches the model to evaluation mode, which disables training-specific behaviors like dropout.

torch.no_grad() tells PyTorch not to track gradients during this pass. This saves memory and speeds up inference since you're only reading predictions, not training.

The loop iterates over every test image, checks if the prediction matches the true label, and tallies correct guesses per class. Finally, overall_acc_cnn stores the overall accuracy so you can compare it against the transfer learning model later.

  • Run the cell by pressing Shift+Enter.
  • You should see a table listing all 10 classes with their individual accuracy percentages, plus an overall accuracy around 65-70%.

Which classes are hardest?

You'll likely notice that vehicle classes (truck, ship, airplane) score higher than animal classes (cat, deer, dog). This makes sense because animals share similar shapes and textures, making them harder for a simple CNN to distinguish.

Seeing 0% for every class?

Make sure Cell 5 (the training cell) ran to completion before running this cell. If training was interrupted, the model weights may not have been updated.

Also check that you're using the same testloader defined in Cell 2. If your runtime disconnected, re-run all cells from the top.

Why is my per-class accuracy showing unexpected results?

Visualize predictions on test images

Numbers tell part of the story. Seeing the actual images alongside their predicted labels makes the model's behavior much more intuitive.

  • Create a new code cell below Cell 6.
  • Paste the following code to display predictions with color-coded labels:
model.eval()
dataiter = iter(testloader)
images, labels = next(dataiter)
images_gpu = images.to(device)

with torch.no_grad():
    outputs = model(images_gpu)
    _, predicted = torch.max(outputs, 1)

fig, axes = plt.subplots(2, 8, figsize=(16, 5))
for i, ax in enumerate(axes.flat):
    img = images[i]
    img = img * torch.tensor([0.2470, 0.2435, 0.2616]).view(3, 1, 1) + torch.tensor([0.4914, 0.4822, 0.4465]).view(3, 1, 1)
    img = img.clamp(0, 1)
    ax.imshow(np.transpose(img.numpy(), (1, 2, 0)))

    pred_label = classes[predicted[i].item()]
    true_label = classes[labels[i].item()]
    color = 'green' if pred_label == true_label else 'red'
    ax.set_title(f"P:{pred_label}\nT:{true_label}", fontsize=8, color=color)
    ax.axis('off')

plt.suptitle("Predictions (Green=Correct, Red=Wrong)", fontsize=12)
plt.tight_layout()
plt.show()

What does this code do?

The code grabs the first batch of test images and runs them through the model. For each image in the 2x8 grid, it prints both the predicted label (P:) and the true label (T:) above the image.

Green text means the model got it right. Red text means it guessed wrong. The unnormalization step reverses the preprocessing so the images display with their original colors.

  • Run the cell by pressing Shift+Enter.
  • You should see a grid of 16 test images with green labels for correct predictions and red labels for incorrect ones.

Images look like solid color blocks?

Check that the unnormalization values match the ones used in Cell 2. The mean and std tensors must be identical. Without proper unnormalization, pixel values stay in a range that displays incorrectly.

Why are my visualization images displaying incorrectly?

✔️ Awesome, I've got everything!

Great. Double check your cells are saved in your notebook before moving on.

ⓧ I'd like to double check the full code

Here is the code for both cells created in this step.

Cell 6:

model.eval()
class_correct = [0] * 10
class_total = [0] * 10

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

print("Per-class accuracy (From-Scratch CNN):")
print("-" * 40)
for i in range(10):
    acc = 100 * class_correct[i] / class_total[i]
    print(f"  {classes[i]:12s}: {acc:.1f}%")

overall_acc_cnn = 100 * sum(class_correct) / sum(class_total)
print(f"\n  Overall: {overall_acc_cnn:.1f}%")

Cell 7:

model.eval()
dataiter = iter(testloader)
images, labels = next(dataiter)
images_gpu = images.to(device)

with torch.no_grad():
    outputs = model(images_gpu)
    _, predicted = torch.max(outputs, 1)

fig, axes = plt.subplots(2, 8, figsize=(16, 5))
for i, ax in enumerate(axes.flat):
    img = images[i]
    img = img * torch.tensor([0.2470, 0.2435, 0.2616]).view(3, 1, 1) + torch.tensor([0.4914, 0.4822, 0.4465]).view(3, 1, 1)
    img = img.clamp(0, 1)
    ax.imshow(np.transpose(img.numpy(), (1, 2, 0)))

    pred_label = classes[predicted[i].item()]
    true_label = classes[labels[i].item()]
    color = 'green' if pred_label == true_label else 'red'
    ax.set_title(f"P:{pred_label}\nT:{true_label}", fontsize=8, color=color)
    ax.axis('off')

plt.suptitle("Predictions (Green=Correct, Red=Wrong)", fontsize=12)
plt.tight_layout()
plt.show()

You now have a clear picture of where your from-scratch CNN excels and where it struggles. Next up, you'll bring in a pretrained ResNet18 model and see how transfer learning can dramatically close those accuracy gaps.

Upgrade to Transfer Learning with ResNet18

Your from-scratch CNN achieved around 65-70% accuracy. That means roughly 1 in 3 images gets misclassified.

Instead of training a model from zero knowledge, what if you started with a model that already understands visual features like edges, textures, and shapes? That is the idea behind transfer learning. In this step, you will load a pretrained ResNet18 model and fine-tune it for CIFAR-10 classification.

In this step, get ready to:
  • Load a pretrained ResNet18 model and adapt it for CIFAR-10.
  • Train the adapted model for 5 epochs.
  • Compare accuracy between the from-scratch CNN and the transfer learning model.
Load pretrained ResNet18 and adapt it for CIFAR-10

ResNet18 is a neural network with 11.7 million parameters that was trained on ImageNet, a dataset of over 1 million images across 1,000 categories. Its lower layers already know how to detect edges, textures, and patterns that are useful for any image task.

  • Create a new code cell below your previous cells.
  • Paste the following code to load ResNet18 and set up training:
from torchvision.models import resnet18, ResNet18_Weights

# Load pretrained ResNet18
resnet_model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)

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

# Training setup
criterion_resnet = nn.CrossEntropyLoss()
optimizer_resnet = optim.SGD(resnet_model.parameters(), lr=0.001, momentum=0.9)

What does this code do?

resnet18(weights=ResNet18_Weights.IMAGENET1K_V1) loads a ResNet18 model with pretrained ImageNet weights. This gives you a model that already understands visual features.

The original ResNet18 outputs 1,000 classes (for ImageNet). resnet_model.fc = nn.Linear(resnet_model.fc.in_features, 10) replaces the final layer so it outputs 10 classes instead, matching CIFAR-10.

The optimizer uses SGD with momentum instead of Adam. SGD with momentum is a common choice for fine-tuning pretrained models because it makes more conservative weight updates.

Why transfer learning instead of a bigger CNN?

You could try building a deeper CNN from scratch, but that requires more data and training time to learn useful features. Transfer learning reuses features that ResNet18 already learned from millions of ImageNet images.

The result is higher accuracy in fewer epochs. This is the same approach used in production ML systems at companies that need to classify custom images without millions of labeled training samples.

Train the fine-tuned model

Because ResNet18 already knows useful features, you only need 5 epochs of training instead of 10. The training loop follows the same pattern as before.

  • In the same cell, paste this training loop below the code you just added:
# Train for 5 epochs
num_epochs_resnet = 5

for epoch in range(num_epochs_resnet):
    resnet_model.train()
    running_loss = 0.0

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

        optimizer_resnet.zero_grad()
        outputs = resnet_model(inputs)
        loss = criterion_resnet(outputs, labels)
        loss.backward()
        optimizer_resnet.step()

        running_loss += loss.item()

    avg_loss = running_loss / len(trainloader)

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

    resnet_acc = 100 * correct / total
    print(f"Epoch [{epoch+1}/{num_epochs_resnet}] Loss: {avg_loss:.4f} | Test Accuracy: {resnet_acc:.2f}%")

print("\nTransfer learning training complete!")

What's happening here?

The training loop is the same pattern you used for the from-scratch CNN: forward pass, compute loss, backpropagate, update weights. The key difference is that most of the model's weights are already good starting values from ImageNet training.

After each epoch, the loop evaluates on the full test set and prints both the loss and accuracy. You should see accuracy climbing much faster than the from-scratch model did.

  • Run the cell by pressing the play button or pressing Shift+Enter.
  • Training takes about 5-6 minutes on a T4 GPU. You should see test accuracy reach approximately 90% or higher by the final epoch.

Seeing much lower accuracy than expected?

Make sure you replaced the final layer BEFORE calling .to(device). If the model is moved to GPU before the layer replacement, the new layer stays on CPU.

Also confirm your runtime still has GPU access. If your Colab session disconnected and reconnected, you may need to re-run all previous cells.

Help me debug low accuracy in my ResNet18 transfer learning model on CIFAR-10.

✔️ Awesome, I've got everything!

Great. Make sure you see accuracy around 90% in the final epoch output before moving on.

ⓧ I'd like to double check the full code

Here is the complete Cell 8 for comparison:

from torchvision.models import resnet18, ResNet18_Weights

# Load pretrained ResNet18
resnet_model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)

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

# Training setup
criterion_resnet = nn.CrossEntropyLoss()
optimizer_resnet = optim.SGD(resnet_model.parameters(), lr=0.001, momentum=0.9)

# Train for 5 epochs
num_epochs_resnet = 5

for epoch in range(num_epochs_resnet):
    resnet_model.train()
    running_loss = 0.0

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

        optimizer_resnet.zero_grad()
        outputs = resnet_model(inputs)
        loss = criterion_resnet(outputs, labels)
        loss.backward()
        optimizer_resnet.step()

        running_loss += loss.item()

    avg_loss = running_loss / len(trainloader)

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

    resnet_acc = 100 * correct / total
    print(f"Epoch [{epoch+1}/{num_epochs_resnet}] Loss: {avg_loss:.4f} | Test Accuracy: {resnet_acc:.2f}%")

print("\nTransfer learning training complete!")
Compare results

Now you have both models trained. Time to put the numbers side by side and see how much transfer learning improved things.

  • Create a new code cell below Cell 8.
  • Paste the following comparison code and run it:
# Final comparison
resnet_model.eval()
correct = 0
total = 0
with torch.no_grad():
    for inputs, labels in testloader:
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = resnet_model(inputs)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

overall_acc_resnet = 100 * correct / total

print("=" * 50)
print("FINAL COMPARISON")
print("=" * 50)
print(f"  From-Scratch CNN (10 epochs):  {overall_acc_cnn:.1f}%")
print(f"  ResNet18 Transfer (5 epochs):  {overall_acc_resnet:.1f}%")
print(f"  Improvement:                   +{overall_acc_resnet - overall_acc_cnn:.1f}%")
print("=" * 50)
print("\nTransfer learning achieves better accuracy in fewer")
print("epochs because ResNet18 already learned useful image")
print("features (edges, textures, shapes) from ImageNet.")

What does this code do?

This cell re-evaluates ResNet18 on the full test set to get a clean final accuracy number. It stores this in overall_acc_resnet so it can print both models' results side by side.

The comparison uses overall_acc_cnn from your earlier evaluation cell. Both numbers come from the same 10,000-image test set, making the comparison fair.

  • You should see a comparison table showing the from-scratch CNN at ~65-70% and ResNet18 at ~90% or higher. The improvement line shows a gain of roughly 20+ percentage points.

Why does transfer learning work so well here?

ResNet18 was trained on over 1 million ImageNet images. During that training, its layers learned to detect edges, textures, shapes, and patterns that generalize across many image tasks.

When you fine-tune it on CIFAR-10, only the final classification layer needs to learn from scratch. All the earlier layers already provide useful feature extraction. This is why 5 epochs beats 10 epochs of the from-scratch CNN.

Seeing a NameError for overall_acc_cnn?

This means the variable from Cell 6 is not in memory. If your Colab runtime disconnected at any point, re-run all cells from Cell 1 through Cell 6 first. Then re-run Cells 8 and 9.

Help me fix a NameError when comparing my CNN and ResNet18 results in Colab.

✔️ Awesome, I've got everything!

Great. Double check you see both accuracy numbers and the improvement line in the output.

ⓧ I'd like to double check the full code

Here is the complete Cell 9 for comparison:

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

overall_acc_resnet = 100 * correct / total

print("=" * 50)
print("FINAL COMPARISON")
print("=" * 50)
print(f"  From-Scratch CNN (10 epochs):  {overall_acc_cnn:.1f}%")
print(f"  ResNet18 Transfer (5 epochs):  {overall_acc_resnet:.1f}%")
print(f"  Improvement:                   +{overall_acc_resnet - overall_acc_cnn:.1f}%")
print("=" * 50)
print("\nTransfer learning achieves better accuracy in fewer")
print("epochs because ResNet18 already learned useful image")
print("features (edges, textures, shapes) from ImageNet.")

Secret mission

Classify Your Own Image

Your ResNet18 model learned to classify CIFAR-10 images with over 90% accuracy. But can it handle real-world photos? Upload your own images and find out where the model shines and where it stumbles.

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 on Google Colab's free tier, so there are no ongoing costs.

Resources you used:

  • Google Colab notebook (CIFAR10_Image_Classifier) saved in Google Drive.
  • CIFAR-10 dataset downloaded inside the Colab runtime.
  • T4 GPU runtime session (free tier, auto-disconnects when idle).

Keep everything running

No action needed. Choose this if you're still actively experimenting with your models or want to try more images.

  • Your notebook is saved in Google Drive automatically.
  • The Colab runtime stays active until it times out from inactivity (usually 90 minutes) or you close the tab.
  • Your trained models, variables, and downloaded data persist as long as the runtime is connected.

Pause - I'll come back to this later

Shut down the runtime to free up the GPU for others. Your notebook code is preserved in Google Drive, but all variables, models, and downloaded data will be lost.

  • In your Colab notebook, click Runtime in the top menu.
  • Select Disconnect and delete runtime.
  • When you return, you will need to re-run all cells from the top to retrain your models and re-download the dataset.

Delete - I don't want to use this again

Remove the notebook from Google Drive entirely. This cannot be undone.

  • In your Colab notebook, click Runtime in the top menu, then select Disconnect and delete runtime. This removes all downloaded data and trained models from the session.
  • Open Google Drive at drive.google.com.
  • Navigate to the Colab Notebooks folder.
  • Right-click on CIFAR10_Image_Classifier.ipynb and select Move to trash.
  • Empty your trash if you want to permanently delete the file.

That's a Wrap!

That's a Wrap!

Nice work! You just built a complete image classification pipeline using PyTorch and Google Colab. You trained two neural networks, compared their performance, and proved firsthand why transfer learning is such a game-changer in deep learning.

You've learned how to:

  • Built a CNN from scratch with convolutional layers, pooling, and fully connected layers, then trained it on 60,000 images to achieve ~65-70% accuracy.
  • Used transfer learning with ResNet18 pretrained on ImageNet to reach ~90%+ accuracy in fewer epochs, demonstrating how pretrained features dramatically boost performance.
  • Evaluated both models with per-class accuracy metrics and color-coded prediction visualizations to understand where each model succeeds and struggles.
  • Secret Mission: Classified your own uploaded images using the trained ResNet18 model with confidence scores and probability breakdowns across all 10 CIFAR-10 categories.

Ready to quiz yourself?