Build a CIFAR-10 CNN with PyTorch

Introduction

30 Second Summary

A computer can get better at naming pictures it has already seen. New pictures reveal whether it learned the patterns or memorized the examples.

In this project, you will use PyTorch to train two convolutional neural networks on CIFAR-10 images using your Windows laptop's CPU. Their learning curves reveal whether data augmentation plus dropout improves performance on unseen images.

What You'll Build

Your finished experiment opens a labeled sample grid, side-by-side learning curves, and a prediction grid from the model chosen by validation accuracy.

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

• Inspect a labeled sample grid showing one example from every CIFAR-10 class before training starts.
• Train a baseline CNN that prints separate loss and accuracy metrics for each epoch.
• Compare baseline and regularized results through saved learning curves, test predictions, and reusable .pth checkpoints.
• Secret Mission: Isolate augmentation and dropout in an ablation study to discover which technique contributes more on your machine.

Are there any prerequisites?

You need basic Python knowledge plus prior exposure to one machine learning model. Step 1 guides you through installing Python, Visual Studio Code, and the required packages from scratch.

Before We Start

Before the hands-on work begins, commit to comparing a baseline CIFAR-10 classifier with a regularized classifier across 10 labels: airplane, automobile, bird, cat, deer, dog, frog, horse, ship, or truck. Validation metrics will select the preferred model during an intentional CPU experiment because the Radeon RX 6500M is not on AMD's supported Windows ROCm matrix.

Set Up the Local CPU Workspace

Comparable classifiers need a repeatable local toolchain. Matching package versions make the baseline and regularized results easier to trust.

The Radeon RX 6500M is absent from AMD's supported Windows ROCm matrix. A CPU environment keeps this project on a reliable path.

Visual Studio Code gives the project a home. Python 3.13 provides the runtime inside an isolated virtual environment.

In this step, get ready to:

• Prepare Visual Studio Code plus Python 3.13.
• Create the .venv environment inside cnn-cifar10-local.
• Install PyTorch 2.14.0, TorchVision 0.29.0, plus Matplotlib 3.11.1 for CPU execution.

Prepare the Windows tools

Visual Studio Code holds the project files. Windows PowerShell runs the setup commands that prepare Python.

• Press the Windows key to open Windows search.
• Type Visual Studio Code.

Check whether Visual Studio Code appears in the search results.

✔️ Visual Studio Code opens

• Press Enter to open Visual Studio Code.

Your editor is ready. Leave it open while you confirm the Python runtime.

ⓧ Visual Studio Code is not installed

The Windows download page contains several installer types. The User Installer is the recommended setup for most learners.

The installer asks you to accept its licence before showing the installation options.

• Open the official Visual Studio Code download page.
• Select x64 in the User Installer row for Windows.
• Open the downloaded installer.
• Accept the licence terms.
• Keep the default installation options.
• Select Install.
• Select Finish to open Visual Studio Code.

Visual Studio Code opens after setup. The editor is now ready for your local project.

Visual Studio Code not opening?

Confirm that you downloaded the Windows x64 User Installer. Run the downloaded installer again if setup closed before completion.

Ask for help with your specific installation symptoms: help me install Visual Studio Code User Setup on Windows.

Python Install Manager provides the supported runtime when Python 3.13 is missing. First, check what Windows currently finds.

• Press the Windows key to open Windows search.
• Type PowerShell.

You should see Windows PowerShell in the search results.

• Press Enter to open Windows PowerShell.
• Check the available Python version by running this command:

python --version

What does this command do?

The --version option prints the Python runtime selected by Windows. The result determines whether you can continue or need the Python Install Manager.

✔️ I see Python 3.13

Your Python runtime is ready. The output begins with Python 3.13, so you can continue to the project environment.

ⓧ I see another Python version

A different runtime could resolve packages differently. Install the supported Python 3.13 runtime before creating the environment.

Windows may ask you to approve package source agreements during the first command. Approve those prompts so installation can continue.

• Install Python Install Manager plus the Python 3.13 runtime by running these commands:

winget install 9NQ7512CXL7T
py install 3.13

What do these commands install?

• The first command installs the official Python Install Manager through Windows Package Manager.
• The second command installs the supported Python 3.13 runtime.

• Confirm that Windows now selects Python 3.13 by running:

python --version

What confirms the runtime?

The command prints the active Python version. An output beginning with Python 3.13 confirms the supported runtime is available.

You should now see an output beginning with Python 3.13.

Still seeing another version?

Close PowerShell after installation. Open a fresh PowerShell window so Windows reloads the available runtime commands.

Ask for help if the older runtime still appears: help me make Python 3.13 available in PowerShell.

ⓧ Python is not found

Windows cannot find a Python runtime yet. Python Install Manager supplies the runtime required by this project.

Windows may ask you to approve package source agreements during the first command. Approve those prompts so installation can continue.

• Install Python Install Manager plus the Python 3.13 runtime by running these commands:

winget install 9NQ7512CXL7T
py install 3.13

What do these commands install?

• The first command installs the official Python Install Manager through Windows Package Manager.
• The second command installs the supported Python 3.13 runtime.

• Confirm that Windows can now find Python by running:

python --version

What confirms the installation?

The command prints the active Python version. An output beginning with Python 3.13 confirms that installation worked.

You should now see an output beginning with Python 3.13.

Python still not found?

Close PowerShell after installation. Open a fresh PowerShell window so Windows reloads the available runtime commands.

Ask for help if Python remains unavailable: help me make Python available in PowerShell.

Isolate the project packages

Your Python runtime is ready. The next layer is a dedicated environment that prevents this project's package versions from mixing with other Python projects.

• Return to Visual Studio Code from earlier.
• Click File in the top menu bar.
• Click Open Folder.

The Windows folder picker opens. You can create the project folder directly from this window.

• Select Desktop in the folder picker.
• Click New folder.
• Enter cnn-cifar10-local as the folder name.
• Press Enter to create the folder.

You should see cnn-cifar10-local listed on your Desktop.

• Select the cnn-cifar10-local folder.
• Click Select Folder.
• Approve the workspace trust prompt if Visual Studio Code shows one.

The Explorer sidebar now shows cnn-cifar10-local as the open folder.

• Click Terminal in the top menu bar.
• Click New Terminal.

The integrated terminal opens with a PowerShell prompt inside cnn-cifar10-local.

• Create the .venv environment plus activate it by running these commands:

python -m venv .venv
.\.venv\Scripts\Activate.ps1

What do these commands do?

• The first command creates a Python environment at cnn-cifar10-local/.venv.
• The second command activates that environment in the current PowerShell session.
• Package commands now target the isolated environment instead of the system runtime.

Your PowerShell prompt should now begin with (.venv). That visible prefix confirms activation.

PowerShell blocking activation?

Confirm that the terminal path ends in cnn-cifar10-local. The activation path only works from the folder that contains .venv.

If PowerShell reports that scripts are disabled, ask for help with the policy on your Windows account: help me activate this .venv safely.

Visual Studio Code also needs to use the same interpreter. Selecting .venv keeps editor features aligned with the active terminal.

• Press Ctrl+Shift+P to open the Command Palette.
• Type Python: Select Interpreter.

You should see Python: Select Interpreter in the command results.

• Select Python: Select Interpreter.
• Select the interpreter whose path ends with cnn-cifar10-local\.venv\Scripts\python.exe.

The Visual Studio Code status bar should now identify Python 3.13 from .venv.

Can't find the .venv interpreter?

Confirm that the .venv folder appears in the Explorer sidebar. Reopen the Command Palette after the environment finishes creating.

If the interpreter command is missing, install the Python extension from Microsoft in the Visual Studio Code Extensions view.

Ask for help if the environment still does not appear: help me connect .venv to Visual Studio Code.

Prepare the CPU package set

The active .venv controls where each package is installed. The CPU wheel source keeps PyTorch aligned with the hardware path chosen for this project.

The PyTorch download can take several minutes on a typical home connection. A pause in the progress output can happen while the larger package files download.

• Install the CPU builds of PyTorch 2.14.0 plus TorchVision 0.29.0 by running:

pip3 install torch==2.14.0 torchvision==0.29.0 --index-url https://download.pytorch.org/whl/cpu

What does this install?

• The torch==2.14.0 pin installs the project framework version.
• The torchvision==0.29.0 pin installs the compatible image dataset tools.
• The CPU package index avoids installing a build intended for unsupported GPU acceleration.

Good progress. The CPU builds of PyTorch plus TorchVision now live inside .venv when the command returns to the PowerShell prompt without an installation error.

PyTorch packages not installing?

Confirm that your prompt begins with (.venv) before retrying. A network interruption can also stop the package download.

Ask for help with the terminal output you received: help me install the pinned PyTorch CPU packages.

• Install Matplotlib 3.11.1 into the same environment by running:

pip install matplotlib==3.11.1

What does this install?

Matplotlib creates the sample grids plus learning curves used later in the project. Pinning 3.11.1 gives both classifier runs the same plotting environment.

The command returns to the PowerShell prompt after Matplotlib and its required packages finish installing.

Matplotlib installation failing?

Check that the terminal prompt still begins with (.venv). An inactive environment can send the installation to another Python runtime.

Ask for help with the package output: help me install Matplotlib in this .venv.

Before you run the final check, do you expect PyTorch to report an available CUDA device on this CPU setup?

• Verify the three package versions plus the CPU execution path by running:

python -c "import torch, torchvision, matplotlib; print(torch.__version__, torchvision.__version__, matplotlib.__version__, torch.cuda.is_available())"

What does this verification prove?

• The first three values report the installed framework and plotting versions.
• The final expression checks whether CUDA acceleration is available to PyTorch.
• A final value of False confirms that this project follows the intended CPU path.

You should see three installed version strings on one line. The line ends with False.

Your local toolchain is ready. Visual Studio Code now uses the isolated environment that holds every pinned CPU package.

Workspace checkpoint

Created Python 3.13 virtual environment with torch==2.14.0 CPU, torchvision==0.29.0, and matplotlib==3.11.1 installed.

Visual Studio Code uses cnn-cifar10-local/.venv as its selected Python interpreter.

Verification command not working?

Confirm that the PowerShell prompt begins with (.venv). Missing package imports usually mean the command is using a different interpreter.

Ask for help with the complete output: help me verify my PyTorch CPU environment.

The CPU workspace is locked in. Next, you'll download CIFAR-10 and turn its 10 classes into a labeled image grid.

Inspect the CIFAR-10 Data

Your local PyTorch CPU workspace is ready. A classifier only becomes meaningful when you can see its inputs and the labels it must predict.

This step uses TorchVision to download CIFAR-10. You will save a labeled grid that makes all ten target classes visible before training begins.

In this step, get ready to:

• Create a runnable data-stage scaffold in cnn_cifar10.py.
• Configure the class labels and normalized image transform.
• Download the dataset and inspect one image from every class.

Build a runnable data-stage scaffold

The program needs a stage argument so each part of the project can run independently. Its first runnable scaffold also creates the local output directory and confirms that computation stays on the CPU.

• Click the New File icon in the VS Code Explorer sidebar.
• Type cnn_cifar10.py as the file name.
• Press Enter to create the file.
• Add the runnable stage scaffold by pasting this code into cnn_cifar10.py:

import argparse
from pathlib import Path

import torch

DEVICE = torch.device("cpu")
DATA_DIR = Path("data")
ARTIFACT_DIR = Path("artifacts")

def parse_args():
    parser = argparse.ArgumentParser(
        description="Train and compare compact CIFAR-10 CNNs on a Windows CPU."
    )
    parser.add_argument(
        "--stage",
        choices=("data", "baseline", "compare", "experiment"),
        default="compare",
    )
    return parser.parse_args()

def main():
    args = parse_args()
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    print(f"Device: {DEVICE}")

    if args.stage == "data":
        return

if __name__ == "__main__":
    main()

What Does This Scaffold Do?

• The argument parser lets Python read the requested program stage.
• The ARTIFACT_DIR path gives every saved chart one consistent destination.
• The DEVICE value keeps this project on the supported CPU path.

• Save cnn_cifar10.py by pressing Ctrl+S.
• Confirm that cnn_cifar10.py appears in the VS Code Explorer sidebar.

Before you run this, picture the result you expect from a data stage that currently returns after confirming its device.

• Run the scaffold in the activated PowerShell terminal with this command:

python cnn_cifar10.py --stage data

What Does This Command Confirm?

The --stage data argument selects the dataset path through the program. The scaffold also creates the artifacts directory for saved plots.

You should see Device: cpu in PowerShell. That is your first working program result, and the chart output folder is now ready.

Script Did Not Print the Device?

Check that your PowerShell prompt still shows the activated .venv environment. Confirm that the terminal is inside the cnn-cifar10-local folder.

Make sure the file name is exactly cnn_cifar10.py.

Help me run the data-stage scaffold from my activated Windows virtual environment.

Define the CIFAR-10 view

Each CIFAR-10 image has three colour channels across a 32-by-32 pixel grid. The evaluation transform converts each image into a tensor and normalizes its channel values for the model.

• Find this import near the top of cnn_cifar10.py:

import torch

Why Locate This Import?

This line marks the point where the dataset and plotting imports belong. Keeping related imports together makes the file easier to compare with the completed program.

• Replace that line with this import group:

import matplotlib.pyplot as plt
import torch
from torchvision import datasets
from torchvision.transforms import v2

What Do These Imports Provide?

• Matplotlib builds and saves the image grid.
• The datasets module downloads and loads CIFAR-10.
• The v2 module supplies the image conversion and normalization steps.

• Find the current configuration block below the imports:

DEVICE = torch.device("cpu")
DATA_DIR = Path("data")
ARTIFACT_DIR = Path("artifacts")

What Does This Configuration Hold?

These paths tell the program where to download data and where to save results. The device value keeps every tensor operation on the CPU.

• Replace the current configuration block with the complete data configuration:

SEED = 42
DEVICE = torch.device("cpu")
DATA_DIR = Path("data")
ARTIFACT_DIR = Path("artifacts")
CLASSES = (
    "airplane",
    "automobile",
    "bird",
    "cat",
    "deer",
    "dog",
    "frog",
    "horse",
    "ship",
    "truck",
)

EVAL_TRANSFORM = v2.Compose(
    [
        v2.ToImage(),
        v2.ToDtype(torch.float32, scale=True),
        v2.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
)

How Is the Data Prepared?

• The CLASSES tuple maps each numeric label to its readable category name.
• The EVAL_TRANSFORM pipeline converts images to floating-point tensors.
• Normalization places all three colour channels on a consistent scale.
• The SEED value gives later dataset splits a repeatable starting point.

• Save cnn_cifar10.py.
• Check the new imports and transform configuration by running the data stage again:

python cnn_cifar10.py --stage data

What Does This Check Prove?

You should see Device: cpu again. This confirms that the dataset modules and normalization pipeline load without stopping the script.

Seeing an Import or Syntax Problem?

Confirm that the selected interpreter still points to cnn-cifar10-local/.venv. Check that every class name remains inside the CLASSES parentheses.

Compare the brackets around EVAL_TRANSFORM carefully. Its transform steps sit inside square brackets within v2.Compose().

Help me diagnose the imports or transform configuration in my CIFAR-10 script.

Download and inspect the samples

Normalized tensors are suitable for model input, but their values need converting back into visible colours for the plot. The sample-grid function collects the first example of each label and arranges the ten images across two rows.

• Insert the image-display functions above parse_args() in cnn_cifar10.py by pasting this code:

def display_image(image):
    return (image * 0.5 + 0.5).clamp(0, 1).permute(1, 2, 0).numpy()

def show_sample_grid():
    dataset = datasets.CIFAR10(
        root=DATA_DIR,
        train=True,
        download=True,
        transform=EVAL_TRANSFORM,
    )
    examples = {}
    for image, label in dataset:
        if label not in examples:
            examples[label] = image
        if len(examples) == len(CLASSES):
            break

    figure, axes = plt.subplots(2, 5, figsize=(11, 5))
    for label, axis in enumerate(axes.flat):
        axis.imshow(display_image(examples[label]))
        axis.set_title(CLASSES[label])
        axis.axis("off")
    figure.suptitle("CIFAR-10 sample classes")
    figure.tight_layout()
    output_path = ARTIFACT_DIR / "cifar10_samples.png"
    figure.savefig(output_path, dpi=150)
    print(f"Saved {output_path}")
    plt.show()
    plt.close(figure)

How Does the Grid Work?

• The display_image() function reverses the normalization and moves the colour channels into the order expected by the plot.
• The examples dictionary keeps one image for each numeric label.
• The two-row figure uses CLASSES to place a readable category name above every image.
• The completed figure is saved to artifacts/cifar10_samples.png at 150 DPI.

• Find the current data-stage branch near the bottom of cnn_cifar10.py:

    if args.stage == "data":
        return

What Does the Current Branch Do?

The stage branch currently exits after the device check. It needs one function call to produce the dataset grid.

• Replace the current data-stage branch with this completed version:

    if args.stage == "data":
        show_sample_grid()
        return

Why Add the Function Call Here?

The branch now runs show_sample_grid() only when the data stage is selected. The return statement stops the program after the visualization finishes.

• Save cnn_cifar10.py.

Before you run this, picture the ten labels you expect to see across the grid.

• Download CIFAR-10 and display its sample grid by running:

python cnn_cifar10.py --stage data

What Happens on This Run?

TorchVision downloads the CIFAR-10 files into data. The program then searches the training data until it has one example from every class.

The plot opens in a separate window and the saved image goes to artifacts/cifar10_samples.png.

You should see a 2-by-5 grid titled CIFAR-10 sample classes. The grid should include airplane, automobile, bird, cat, deer, dog, frog, horse, ship, and truck.

• Close the plot window after inspecting all ten labels.
• Expand the artifacts folder in the VS Code Explorer sidebar.
• Confirm that cifar10_samples.png appears inside the artifacts folder.
• Expand the data folder to confirm that the CIFAR-10 files were downloaded locally.

Great result. You can now see the ten categories your CNN must learn to separate, and the sample chart is saved for later comparison.

✔️ Awesome, I've got everything!

Your dataset is downloaded and your labeled sample grid is saved. Double check that cnn_cifar10.py is saved before moving on.

ⓧ I'd like to double check the full code

Compare your complete cnn_cifar10.py file with this version.

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import torch
from torchvision import datasets
from torchvision.transforms import v2

SEED = 42
DEVICE = torch.device("cpu")
DATA_DIR = Path("data")
ARTIFACT_DIR = Path("artifacts")
CLASSES = (
    "airplane",
    "automobile",
    "bird",
    "cat",
    "deer",
    "dog",
    "frog",
    "horse",
    "ship",
    "truck",
)

EVAL_TRANSFORM = v2.Compose(
    [
        v2.ToImage(),
        v2.ToDtype(torch.float32, scale=True),
        v2.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
)


def display_image(image):
    return (image * 0.5 + 0.5).clamp(0, 1).permute(1, 2, 0).numpy()


def show_sample_grid():
    dataset = datasets.CIFAR10(
        root=DATA_DIR,
        train=True,
        download=True,
        transform=EVAL_TRANSFORM,
    )
    examples = {}
    for image, label in dataset:
        if label not in examples:
            examples[label] = image
        if len(examples) == len(CLASSES):
            break

    figure, axes = plt.subplots(2, 5, figsize=(11, 5))
    for label, axis in enumerate(axes.flat):
        axis.imshow(display_image(examples[label]))
        axis.set_title(CLASSES[label])
        axis.axis("off")
    figure.suptitle("CIFAR-10 sample classes")
    figure.tight_layout()
    output_path = ARTIFACT_DIR / "cifar10_samples.png"
    figure.savefig(output_path, dpi=150)
    print(f"Saved {output_path}")
    plt.show()
    plt.close(figure)


def parse_args():
    parser = argparse.ArgumentParser(
        description="Train and compare compact CIFAR-10 CNNs on a Windows CPU."
    )
    parser.add_argument(
        "--stage",
        choices=("data", "baseline", "compare", "experiment"),
        default="compare",
    )
    return parser.parse_args()


def main():
    args = parse_args()
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    print(f"Device: {DEVICE}")

    if args.stage == "data":
        show_sample_grid()
        return


if __name__ == "__main__":
    main()

What Should Match?

This file combines the CPU configuration, ten class labels, normalized evaluation transform, sample-grid functions, data-stage argument, and main entry point. Matching these pieces produces the downloaded dataset and saved grid you just inspected.

Plot Missing or Labels Incomplete?

Check that show_sample_grid() is indented inside the data-stage branch. Confirm that all ten names remain inside CLASSES.

If the download stopped, keep the activated PowerShell terminal open and run the data-stage command again. TorchVision reuses files that have already completed downloading.

Help me diagnose a missing CIFAR-10 grid or an incomplete set of labels.

Your classifier now has a visible problem to solve. Next, you will build the baseline CNN and measure whether its training progress carries over to validation data.

Train the Baseline CNN

Your sample grid proved that the program can load CIFAR-10 images with the correct labels. Now your PyTorch program needs a convolutional neural network that can learn patterns from those images.

Training accuracy alone can hide poor performance on unfamiliar images. Separate validation measurements expose that blind spot before the final test set is used.

In this step, get ready to:

• Build the baseline CNN with convolutional layers and a linear classifier.
• Train the model while measuring separate training and validation results.
• Save the best checkpoint and inspect its learning curves and test predictions.

Build the network and data loaders

The network turns each image into logits for the ten classes. The loaders keep training, validation, and test images separate throughout the experiment.

• Return to cnn_cifar10.py in the Visual Studio Code tab from earlier.
• Replace the import section at the top of the file with this expanded import group:

import argparse
import copy
from pathlib import Path

import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets
from torchvision.transforms import v2

What do these imports add?

• copy preserves the model state with the highest validation accuracy.
• torch.nn provides the layers and loss function used by the CNN.
• DataLoader creates shuffled training batches and stable evaluation batches.
• Subset applies the deterministic training, validation, and test indices.

• Insert the network class immediately below EVAL_TRANSFORM by pasting this code:

class SmallCNN(nn.Module):
    def __init__(self, dropout_probability):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Dropout(p=dropout_probability),
            nn.Linear(128 * 4 * 4, 128),
            nn.ReLU(),
            nn.Dropout(p=dropout_probability),
            nn.Linear(128, len(CLASSES)),
        )

    def forward(self, images):
        return self.classifier(self.features(images))

How does the CNN process an image?

• Conv2d learns local patterns such as edges and textures.
• ReLU keeps useful positive activations between layers.
• MaxPool2d reduces each spatial dimension from 32 to 16 to 8 to 4.
• Flatten converts the final feature maps into values for the linear classifier.
• SmallCNN.forward() returns one class score for every name in CLASSES.

• Save cnn_cifar10.py.
• Check the updated imports and network syntax by running this command in the PowerShell terminal from earlier:

python -m py_compile cnn_cifar10.py

What does this check prove?

The command compiles the source file into Python bytecode. A successful check returns to the PowerShell prompt without a syntax error.

Seeing a syntax error?

Check that the import group remains at the top of cnn_cifar10.py. Make sure SmallCNN sits below the closing parenthesis of EVAL_TRANSFORM.

Use the reported line number to find a missing parenthesis or incorrect indentation. Help me fix the syntax error in my CNN imports or SmallCNN class.

• Insert the transform selector immediately below SmallCNN by pasting this function:

def transform_for(variant):
    if variant in {"augment", "improved"}:
        return AUGMENTED_TRANSFORM
    return EVAL_TRANSFORM

What does the transform selector do?

transform_for() assigns the unchanged evaluation transform to the baseline. This keeps the first experiment focused on learning from one consistent form of every training image.

• Insert the dropout selector immediately below transform_for() by pasting this function:

def dropout_for(variant):
    if variant in {"dropout", "improved"}:
        return 0.4
    return 0.0

How is the baseline controlled?

dropout_for() returns 0.0 for the baseline. This gives you a controlled model to compare with regularized variants later.

• Save cnn_cifar10.py.
• Confirm both selector functions compile by running:

python -m py_compile cnn_cifar10.py

What should happen?

PowerShell should return to the prompt without printing an error. Python can compile these functions before every possible branch is used.

Are the selectors failing to compile?

Check that both function definitions begin at the left edge of the file. Keep each return statement indented inside its matching function.

Help me correct the indentation in transform_for() and dropout_for().

A deterministic split gives both experiments the same training and validation images. The fixed SEED makes that comparison reproducible.

• Insert the first part of make_loaders() immediately below dropout_for() by pasting this code:

def make_loaders(variant, train_size, val_size, test_size, batch_size):
    training_dataset = datasets.CIFAR10(
        root=DATA_DIR,
        train=True,
        download=True,
        transform=transform_for(variant),
    )
    validation_dataset = datasets.CIFAR10(
        root=DATA_DIR,
        train=True,
        download=True,
        transform=EVAL_TRANSFORM,
    )
    test_dataset = datasets.CIFAR10(
        root=DATA_DIR,
        train=False,
        download=True,
        transform=EVAL_TRANSFORM,
    )

    if train_size + val_size > len(training_dataset):
        raise ValueError("train-size plus val-size exceeds the training dataset")

    generator = torch.Generator().manual_seed(SEED)
    shuffled_indices = torch.randperm(
        len(training_dataset), generator=generator
    ).tolist()
    train_indices = shuffled_indices[:train_size]
    val_indices = shuffled_indices[train_size : train_size + val_size]
    test_indices = list(range(min(test_size, len(test_dataset))))

How are the datasets separated?

• training_dataset provides the images used to update the model.
• validation_dataset measures progress without updating the model.
• test_dataset provides the final held-out measurement.
• torch.randperm() shuffles the training indices with the fixed generator seed.

• Continue make_loaders() directly below test_indices by pasting this indented code:

    train_loader = DataLoader(
        Subset(training_dataset, train_indices),
        batch_size=batch_size,
        shuffle=True,
        num_workers=0,
    )
    val_loader = DataLoader(
        Subset(validation_dataset, val_indices),
        batch_size=batch_size,
        shuffle=False,
        num_workers=0,
    )
    test_loader = DataLoader(
        Subset(test_dataset, test_indices),
        batch_size=batch_size,
        shuffle=False,
        num_workers=0,
    )
    return train_loader, val_loader, test_loader

Why use three loaders?

The training loader shuffles batches so the model sees samples in changing orders. The validation and test loaders preserve their order because evaluation does not need random sampling.

num_workers=0 keeps dataset loading in the main process. This is the stable Windows configuration for this local experiment.

• Save cnn_cifar10.py.
• Confirm the completed loader function compiles by running:

python -m py_compile cnn_cifar10.py

What does success look like?

PowerShell should return to the prompt without an error. Your file now contains a valid CNN architecture and deterministic data-loading pipeline.

Seeing an indentation error in make_loaders?

Keep every loader definition indented by four spaces inside make_loaders(). Place the return statement at the same indentation level as the three loader definitions.

Help me repair the make_loaders() indentation.

Add training and checkpointing

Each pass needs a different mode. Training calculates gradients and updates weights. Evaluation measures performance without changing the model.

• Insert the first part of run_pass() immediately below make_loaders() by pasting this code:

def run_pass(model, loader, loss_function, optimizer=None):
    training = optimizer is not None
    if training:
        model.train()
    else:
        model.eval()

    total_loss = 0.0
    correct = 0
    total = 0

    with torch.set_grad_enabled(training):
        for images, labels in loader:
            images = images.to(DEVICE)
            labels = labels.to(DEVICE)

            if training:
                optimizer.zero_grad(set_to_none=True)

            logits = model(images)
            loss = loss_function(logits, labels)

            if training:
                loss.backward()
                optimizer.step()

How does a training pass work?

• model.train() places the network in training mode when an optimizer is present.
• model.eval() places the network in evaluation mode for validation or testing.
• loss.backward() calculates how each model parameter contributed to the loss.
• optimizer.step() updates the parameters after those gradients are calculated.

• Continue run_pass() directly below optimizer.step() by pasting this indented code:

            predictions = logits.argmax(dim=1)
            total_loss += loss.item() * labels.size(0)
            correct += (predictions == labels).sum().item()
            total += labels.size(0)

    return total_loss / total, correct / total

How are the metrics calculated?

argmax() chooses the class with the highest logit for each image. The function returns average loss and the fraction of correct predictions.

• Save cnn_cifar10.py.
• Check the shared training and evaluation pass by running:

python -m py_compile cnn_cifar10.py

What has this check confirmed?

The completed pass compiles without a syntax error. The same metric calculation can now be reused for training, validation, and testing.

Does run_pass fail to compile?

Keep the prediction and metric lines inside the for loop. Keep the final return statement outside that loop while leaving it inside run_pass().

Help me fix the nesting in run_pass().

Cross-entropy loss scores how well the logits match the true labels. The Adam optimizer uses that signal to update the network after every training batch.

• Insert the first part of train_variant() immediately below run_pass() by pasting this code:

def train_variant(variant, train_loader, val_loader, epochs):
    torch.manual_seed(SEED)
    model = SmallCNN(dropout_for(variant)).to(DEVICE)
    loss_function = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    history = {"train_loss": [], "val_loss": [], "train_acc": [], "val_acc": []}
    best_validation_accuracy = -1.0
    best_state = None

    for epoch in range(1, epochs + 1):
        train_loss, train_accuracy = run_pass(
            model, train_loader, loss_function, optimizer
        )
        val_loss, val_accuracy = run_pass(model, val_loader, loss_function)

        history["train_loss"].append(train_loss)
        history["val_loss"].append(val_loss)
        history["train_acc"].append(train_accuracy)
        history["val_acc"].append(val_accuracy)

        print(
            f"{variant:>9} | epoch {epoch}/{epochs} | "
            f"train loss {train_loss:.4f} | val loss {val_loss:.4f} | "
            f"train acc {train_accuracy:.3f} | val acc {val_accuracy:.3f}"
        )

What does the epoch loop record?

• history stores separate training and validation loss values.
• history also stores separate training and validation accuracy values.
• lr=0.001 sets the learning rate used by Adam.
• print() exposes all four measurements after each epoch.

• Complete train_variant() directly below the epoch output by pasting this indented code:

        if val_accuracy > best_validation_accuracy:
            best_validation_accuracy = val_accuracy
            best_state = copy.deepcopy(model.state_dict())

    model.load_state_dict(best_state)
    return model, history

Why restore the best state?

Checkpointing keeps the parameters from the epoch with the highest validation accuracy. The returned model therefore represents the best observed validation result instead of automatically using the final epoch.

• Save cnn_cifar10.py.
• Confirm the training loop and best-state selection compile by running:

python -m py_compile cnn_cifar10.py

What does success confirm?

PowerShell should return to the prompt without an error. The program can now collect epoch history and restore the validation-selected state.

Is best_state causing an error?

Keep the validation comparison inside the epoch loop. Place model.load_state_dict(best_state) after the loop so it runs once training finishes.

Help me fix the best-state checkpointing block.

A saved checkpoint holds the selected model state plus the measurements and configuration that produced it. That context makes later comparisons traceable.

• Insert the first part of load_or_train() immediately below train_variant() by pasting this code:

def load_or_train(variant, args):
    train_loader, val_loader, test_loader = make_loaders(
        variant,
        args.train_size,
        args.val_size,
        args.test_size,
        args.batch_size,
    )
    checkpoint_path = ARTIFACT_DIR / f"{variant}.pth"
    expected_config = {
        "variant": variant,
        "epochs": args.epochs,
        "train_size": args.train_size,
        "val_size": args.val_size,
    }

    if checkpoint_path.exists() and not args.retrain:
        checkpoint = torch.load(
            checkpoint_path, map_location=DEVICE, weights_only=True
        )
        if checkpoint.get("config") == expected_config:
            model = SmallCNN(dropout_for(variant)).to(DEVICE)
            model.load_state_dict(checkpoint["model_state_dict"])
            print(f"Reused {checkpoint_path}")
            return model, checkpoint["history"], test_loader

When is a checkpoint reused?

load_or_train() checks whether the saved configuration matches the current variant, epoch count, training size, and validation size. A matching checkpoint avoids repeating the same training run.

• Complete load_or_train() directly below its checkpoint branch by pasting this code:

    model, history = train_variant(variant, train_loader, val_loader, args.epochs)
    torch.save(
        {
            "model_state_dict": model.state_dict(),
            "history": history,
            "config": expected_config,
        },
        checkpoint_path,
    )
    print(f"Saved {checkpoint_path}")
    return model, history, test_loader


def test_model(model, test_loader):
    loss_function = nn.CrossEntropyLoss()
    return run_pass(model, test_loader, loss_function)

What gets saved and tested?

torch.save() stores the validation-selected parameters, metric history, and experiment configuration. test_model() measures the restored model on the held-out test loader.

• Save cnn_cifar10.py.
• Check the checkpoint and test functions by running:

python -m py_compile cnn_cifar10.py

What is ready now?

The program now compiles with a complete training pipeline. It can train a variant, preserve its strongest validation state, save its experiment record, and calculate held-out test metrics.

Does load_or_train fail to compile?

Keep the training and save block inside load_or_train(). Start test_model() at the left edge after that function ends.

Help me separate load_or_train() from test_model().

Plot and inspect the baseline

Learning curves make overfitting visible by placing training and validation behavior side by side. A prediction grid then connects the aggregate test score to individual images.

• Insert plot_history() immediately below show_sample_grid() by pasting this function:

def plot_history(history, title, output_name):
    epochs = range(1, len(history["train_loss"]) + 1)
    figure, axes = plt.subplots(1, 2, figsize=(11, 4))
    axes[0].plot(epochs, history["train_loss"], marker="o", label="train")
    axes[0].plot(epochs, history["val_loss"], marker="o", label="validation")
    axes[0].set_title("Loss")
    axes[0].set_xlabel("Epoch")
    axes[0].legend()

    axes[1].plot(epochs, history["train_acc"], marker="o", label="train")
    axes[1].plot(epochs, history["val_acc"], marker="o", label="validation")
    axes[1].set_title("Accuracy")
    axes[1].set_xlabel("Epoch")
    axes[1].legend()

    figure.suptitle(title)
    figure.tight_layout()
    output_path = ARTIFACT_DIR / output_name
    figure.savefig(output_path, dpi=150)
    print(f"Saved {output_path}")
    plt.show()
    plt.close(figure)

What do the learning curves reveal?

The left chart compares training loss with validation loss. The right chart compares training accuracy with validation accuracy.

A growing separation between the training and validation curves shows that gains on the training subset are not transferring equally to held-out images.

• Insert show_predictions() immediately below plot_history() by pasting this function:

def show_predictions(model, test_loader, title, output_name):
    model.eval()
    images, labels = next(iter(test_loader))
    with torch.no_grad():
        predictions = model(images.to(DEVICE)).argmax(dim=1).cpu()

    figure, axes = plt.subplots(2, 4, figsize=(10, 5))
    for index, axis in enumerate(axes.flat):
        predicted_name = CLASSES[predictions[index].item()]
        actual_name = CLASSES[labels[index].item()]
        axis.imshow(display_image(images[index]))
        axis.set_title(f"P: {predicted_name}\nA: {actual_name}")
        axis.axis("off")

    figure.suptitle(title)
    figure.tight_layout()
    output_path = ARTIFACT_DIR / output_name
    figure.savefig(output_path, dpi=150)
    print(f"Saved {output_path}")
    plt.show()
    plt.close(figure)

How are the predictions displayed?

torch.no_grad() disables gradient tracking because prediction does not update the model. Each image is labeled with its predicted class and actual class.

• Save cnn_cifar10.py.
• Confirm both plotting functions compile by running:

python -m py_compile cnn_cifar10.py

What has this added?

The program can now turn metric history into a saved learning-curve chart. It can also create a saved grid of labeled predictions from the selected model.

Are the plot functions failing?

Place both functions after the closing line of show_sample_grid(). Keep each plotting function at the left edge of the file.

Help me locate the plotting syntax error.

The command-line arguments set the experiment budget. The defaults use 8,000 training images, 2,000 validation images, 2,000 test images, three epochs, and batches of 64.

• Find the existing parse_args() function near the bottom of cnn_cifar10.py.
• Replace that complete function with this expanded version:

def parse_args():
    parser = argparse.ArgumentParser(
        description="Train and compare compact CIFAR-10 CNNs on a Windows CPU."
    )
    parser.add_argument(
        "--stage",
        choices=("data", "baseline", "compare", "experiment"),
        default="compare",
    )
    parser.add_argument(
        "--variant", choices=("augment", "dropout"), default="augment"
    )
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--train-size", type=int, default=8000)
    parser.add_argument("--val-size", type=int, default=2000)
    parser.add_argument("--test-size", type=int, default=2000)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--retrain", action="store_true")
    return parser.parse_args()

What do the baseline defaults control?

The size arguments control how many images enter each split. The epoch and batch arguments control how long training runs and how many images are processed together.

--retrain provides an explicit way to ignore a matching checkpoint during a future rerun.

• Find the existing main() function below parse_args().
• Replace that complete function with this baseline-capable version:

def main():
    args = parse_args()
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    print(f"Device: {DEVICE}")

    if args.stage == "data":
        show_sample_grid()
        return

    if args.stage == "baseline":
        model, history, test_loader = load_or_train("baseline", args)
        test_loss, test_accuracy = test_model(model, test_loader)
        print(f"baseline | test loss {test_loss:.4f} | test acc {test_accuracy:.3f}")
        plot_history(history, "Baseline CNN", "baseline_curves.png")
        show_predictions(
            model,
            test_loader,
            "Baseline CNN predictions",
            "baseline_predictions.png",
        )
        return

What happens during the baseline stage?

main() trains or loads the baseline before measuring its test performance. It then saves and displays the learning curves before displaying the prediction grid.

The two plot windows open one at a time. Closing the learning-curves window lets the prediction window open next.

• Save cnn_cifar10.py.
• Check the completed baseline stage by running:

python -m py_compile cnn_cifar10.py

What does this final syntax check prove?

The complete baseline program compiles without a syntax error. The next run can now reach the CNN, loaders, training loop, checkpoint, charts, and prediction grid.

Is the baseline stage failing to compile?

Keep the if args.stage == "baseline": block indented inside main(). Keep the final return statement aligned with the other lines inside that block.

Help me fix the baseline branch in main().

CPU training pauses between epoch summaries because the program prints only after a complete epoch. Leave PowerShell running while it processes the batches.

• Prepare for two plot windows by leaving Visual Studio Code and PowerShell open.

Before you run the baseline, do you expect its training and validation curves to stay together or separate over the three epochs?

• Train and evaluate the baseline by running this command:

python cnn_cifar10.py --stage baseline

What should you observe?

The terminal prints the CPU device followed by three epoch summaries. Each summary includes training loss, validation loss, training accuracy, and validation accuracy.

After training, the terminal prints test loss and test accuracy. It also confirms that artifacts/baseline.pth, artifacts/baseline_curves.png, and artifacts/baseline_predictions.png were saved.

• Inspect the training and validation lines in the learning-curves window.
• Close the learning-curves window to continue the program.
• Inspect the predicted and actual labels in the prediction-grid window.
• Close the prediction-grid window to finish the run.

A higher training accuracy than validation accuracy is evidence of a generalization gap. Incorrect prediction labels show where that gap affects individual test images.

Did training or plotting fail?

If the program reports a syntax problem, run the compile check again and use its line number. If a plot window remains open, close it so the next part of the program can continue.

If the checkpoint was created by an incomplete earlier attempt, remove only artifacts/baseline.pth and rerun the baseline command.

Help me diagnose my baseline training run.

✔️ Awesome, I've got everything!

Your baseline is trained. You now have a validation-selected checkpoint, three epochs of metrics, a learning-curves chart, and a labeled test-prediction grid.

ⓧ I'd like to double check the full code

Compare your complete cnn_cifar10.py with this cumulative baseline version.

import argparse
import copy
from pathlib import Path

import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets
from torchvision.transforms import v2

SEED = 42
DEVICE = torch.device("cpu")
DATA_DIR = Path("data")
ARTIFACT_DIR = Path("artifacts")
CLASSES = (
    "airplane",
    "automobile",
    "bird",
    "cat",
    "deer",
    "dog",
    "frog",
    "horse",
    "ship",
    "truck",
)

EVAL_TRANSFORM = v2.Compose(
    [
        v2.ToImage(),
        v2.ToDtype(torch.float32, scale=True),
        v2.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
)


class SmallCNN(nn.Module):
    def __init__(self, dropout_probability):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Dropout(p=dropout_probability),
            nn.Linear(128 * 4 * 4, 128),
            nn.ReLU(),
            nn.Dropout(p=dropout_probability),
            nn.Linear(128, len(CLASSES)),
        )

    def forward(self, images):
        return self.classifier(self.features(images))


def transform_for(variant):
    if variant in {"augment", "improved"}:
        return AUGMENTED_TRANSFORM
    return EVAL_TRANSFORM


def dropout_for(variant):
    if variant in {"dropout", "improved"}:
        return 0.4
    return 0.0


def make_loaders(variant, train_size, val_size, test_size, batch_size):
    training_dataset = datasets.CIFAR10(
        root=DATA_DIR,
        train=True,
        download=True,
        transform=transform_for(variant),
    )
    validation_dataset = datasets.CIFAR10(
        root=DATA_DIR,
        train=True,
        download=True,
        transform=EVAL_TRANSFORM,
    )
    test_dataset = datasets.CIFAR10(
        root=DATA_DIR,
        train=False,
        download=True,
        transform=EVAL_TRANSFORM,
    )

    if train_size + val_size > len(training_dataset):
        raise ValueError("train-size plus val-size exceeds the training dataset")

    generator = torch.Generator().manual_seed(SEED)
    shuffled_indices = torch.randperm(
        len(training_dataset), generator=generator
    ).tolist()
    train_indices = shuffled_indices[:train_size]
    val_indices = shuffled_indices[train_size : train_size + val_size]
    test_indices = list(range(min(test_size, len(test_dataset))))

    train_loader = DataLoader(
        Subset(training_dataset, train_indices),
        batch_size=batch_size,
        shuffle=True,
        num_workers=0,
    )
    val_loader = DataLoader(
        Subset(validation_dataset, val_indices),
        batch_size=batch_size,
        shuffle=False,
        num_workers=0,
    )
    test_loader = DataLoader(
        Subset(test_dataset, test_indices),
        batch_size=batch_size,
        shuffle=False,
        num_workers=0,
    )
    return train_loader, val_loader, test_loader


def run_pass(model, loader, loss_function, optimizer=None):
    training = optimizer is not None
    if training:
        model.train()
    else:
        model.eval()

    total_loss = 0.0
    correct = 0
    total = 0

    with torch.set_grad_enabled(training):
        for images, labels in loader:
            images = images.to(DEVICE)
            labels = labels.to(DEVICE)

            if training:
                optimizer.zero_grad(set_to_none=True)

            logits = model(images)
            loss = loss_function(logits, labels)

            if training:
                loss.backward()
                optimizer.step()

            predictions = logits.argmax(dim=1)
            total_loss += loss.item() * labels.size(0)
            correct += (predictions == labels).sum().item()
            total += labels.size(0)

    return total_loss / total, correct / total


def train_variant(variant, train_loader, val_loader, epochs):
    torch.manual_seed(SEED)
    model = SmallCNN(dropout_for(variant)).to(DEVICE)
    loss_function = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    history = {"train_loss": [], "val_loss": [], "train_acc": [], "val_acc": []}
    best_validation_accuracy = -1.0
    best_state = None

    for epoch in range(1, epochs + 1):
        train_loss, train_accuracy = run_pass(
            model, train_loader, loss_function, optimizer
        )
        val_loss, val_accuracy = run_pass(model, val_loader, loss_function)

        history["train_loss"].append(train_loss)
        history["val_loss"].append(val_loss)
        history["train_acc"].append(train_accuracy)
        history["val_acc"].append(val_accuracy)

        print(
            f"{variant:>9} | epoch {epoch}/{epochs} | "
            f"train loss {train_loss:.4f} | val loss {val_loss:.4f} | "
            f"train acc {train_accuracy:.3f} | val acc {val_accuracy:.3f}"
        )

        if val_accuracy > best_validation_accuracy:
            best_validation_accuracy = val_accuracy
            best_state = copy.deepcopy(model.state_dict())

    model.load_state_dict(best_state)
    return model, history


def load_or_train(variant, args):
    train_loader, val_loader, test_loader = make_loaders(
        variant,
        args.train_size,
        args.val_size,
        args.test_size,
        args.batch_size,
    )
    checkpoint_path = ARTIFACT_DIR / f"{variant}.pth"
    expected_config = {
        "variant": variant,
        "epochs": args.epochs,
        "train_size": args.train_size,
        "val_size": args.val_size,
    }

    if checkpoint_path.exists() and not args.retrain:
        checkpoint = torch.load(
            checkpoint_path, map_location=DEVICE, weights_only=True
        )
        if checkpoint.get("config") == expected_config:
            model = SmallCNN(dropout_for(variant)).to(DEVICE)
            model.load_state_dict(checkpoint["model_state_dict"])
            print(f"Reused {checkpoint_path}")
            return model, checkpoint["history"], test_loader

    model, history = train_variant(variant, train_loader, val_loader, args.epochs)
    torch.save(
        {
            "model_state_dict": model.state_dict(),
            "history": history,
            "config": expected_config,
        },
        checkpoint_path,
    )
    print(f"Saved {checkpoint_path}")
    return model, history, test_loader


def test_model(model, test_loader):
    loss_function = nn.CrossEntropyLoss()
    return run_pass(model, test_loader, loss_function)


def display_image(image):
    return (image * 0.5 + 0.5).clamp(0, 1).permute(1, 2, 0).numpy()


def show_sample_grid():
    dataset = datasets.CIFAR10(
        root=DATA_DIR,
        train=True,
        download=True,
        transform=EVAL_TRANSFORM,
    )
    examples = {}
    for image, label in dataset:
        if label not in examples:
            examples[label] = image
        if len(examples) == len(CLASSES):
            break

    figure, axes = plt.subplots(2, 5, figsize=(11, 5))
    for label, axis in enumerate(axes.flat):
        axis.imshow(display_image(examples[label]))
        axis.set_title(CLASSES[label])
        axis.axis("off")
    figure.suptitle("CIFAR-10 sample classes")
    figure.tight_layout()
    output_path = ARTIFACT_DIR / "cifar10_samples.png"
    figure.savefig(output_path, dpi=150)
    print(f"Saved {output_path}")
    plt.show()
    plt.close(figure)


def plot_history(history, title, output_name):
    epochs = range(1, len(history["train_loss"]) + 1)
    figure, axes = plt.subplots(1, 2, figsize=(11, 4))
    axes[0].plot(epochs, history["train_loss"], marker="o", label="train")
    axes[0].plot(epochs, history["val_loss"], marker="o", label="validation")
    axes[0].set_title("Loss")
    axes[0].set_xlabel("Epoch")
    axes[0].legend()

    axes[1].plot(epochs, history["train_acc"], marker="o", label="train")
    axes[1].plot(epochs, history["val_acc"], marker="o", label="validation")
    axes[1].set_title("Accuracy")
    axes[1].set_xlabel("Epoch")
    axes[1].legend()

    figure.suptitle(title)
    figure.tight_layout()
    output_path = ARTIFACT_DIR / output_name
    figure.savefig(output_path, dpi=150)
    print(f"Saved {output_path}")
    plt.show()
    plt.close(figure)


def show_predictions(model, test_loader, title, output_name):
    model.eval()
    images, labels = next(iter(test_loader))
    with torch.no_grad():
        predictions = model(images.to(DEVICE)).argmax(dim=1).cpu()

    figure, axes = plt.subplots(2, 4, figsize=(10, 5))
    for index, axis in enumerate(axes.flat):
        predicted_name = CLASSES[predictions[index].item()]
        actual_name = CLASSES[labels[index].item()]
        axis.imshow(display_image(images[index]))
        axis.set_title(f"P: {predicted_name}\nA: {actual_name}")
        axis.axis("off")

    figure.suptitle(title)
    figure.tight_layout()
    output_path = ARTIFACT_DIR / output_name
    figure.savefig(output_path, dpi=150)
    print(f"Saved {output_path}")
    plt.show()
    plt.close(figure)


def parse_args():
    parser = argparse.ArgumentParser(
        description="Train and compare compact CIFAR-10 CNNs on a Windows CPU."
    )
    parser.add_argument(
        "--stage",
        choices=("data", "baseline", "compare", "experiment"),
        default="compare",
    )
    parser.add_argument(
        "--variant", choices=("augment", "dropout"), default="augment"
    )
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--train-size", type=int, default=8000)
    parser.add_argument("--val-size", type=int, default=2000)
    parser.add_argument("--test-size", type=int, default=2000)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--retrain", action="store_true")
    return parser.parse_args()


def main():
    args = parse_args()
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    print(f"Device: {DEVICE}")

    if args.stage == "data":
        show_sample_grid()
        return

    if args.stage == "baseline":
        model, history, test_loader = load_or_train("baseline", args)
        test_loss, test_accuracy = test_model(model, test_loader)
        print(f"baseline | test loss {test_loss:.4f} | test acc {test_accuracy:.3f}")
        plot_history(history, "Baseline CNN", "baseline_curves.png")
        show_predictions(
            model,
            test_loader,
            "Baseline CNN predictions",
            "baseline_predictions.png",
        )
        return


if __name__ == "__main__":
    main()

That is the baseline complete. Next, you will add augmentation and dropout to test whether regularization improves validation behavior.

Regularize and Compare the Models

Your baseline CNN now reports separate training and validation results. Those curves reveal whether the model is learning patterns that carry beyond its training images.

This step tests data augmentation with random crops and horizontal flips. It also adds dropout to reduce dependence on individual activations. You will compare both models using validation measurements before viewing the selected model's test predictions.

In this step, get ready to:

• Add augmented training images and dropout to the improved model.
• Reuse the saved baseline checkpoint while training the improved model under the same experiment settings.
• Compare validation histories and display predictions from the validation-selected model.

Add augmentation and dropout

The improved variant needs varied training inputs while keeping validation and test images consistent. Random crops shift the visible framing. Horizontal flips provide another plausible view of an object.

• In the VS Code Explorer sidebar, return to cnn_cifar10.py.
• Locate the closing parenthesis of EVAL_TRANSFORM.
• Add the augmented transform immediately below it by pasting this code:

AUGMENTED_TRANSFORM = v2.Compose(
    [
        v2.ToImage(),
        v2.RandomCrop(32, padding=4),
        v2.RandomHorizontalFlip(p=0.5),
        v2.ToDtype(torch.float32, scale=True),
        v2.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
)

How does augmentation help?

• The random crop changes the framing while preserving the image size expected by the network.
• The horizontal flip gives the model another valid orientation with a probability of 0.5.
• The final two transforms keep the same data type and normalization used by evaluation images.

• Locate the existing transform_for() function.
• Replace that function with this variant-aware version:

def transform_for(variant):
    if variant in {"augment", "improved"}:
        return AUGMENTED_TRANSFORM
    return EVAL_TRANSFORM

What does this function control?

The augment and improved variants receive randomized training images. Every other variant keeps the normalized evaluation transform.

• Locate the existing dropout_for() function.
• Replace that function with this variant-aware version:

def dropout_for(variant):
    if variant in {"dropout", "improved"}:
        return 0.4
    return 0.0

How is dropout isolated?

The dropout and improved variants use a dropout probability of 0.4. The baseline stays at 0.0 so its behavior remains unchanged.

Seeing a problem in the transform functions?

Check that AUGMENTED_TRANSFORM sits above transform_for(). Confirm that both variant names use lowercase letters inside quotation marks.

Ask for help checking the regularization code in cnn_cifar10.py.

Build the comparison path

A fair comparison keeps the data split and training budget fixed. The comparison chart therefore focuses on validation loss and validation accuracy from the two saved histories.

• Locate plot_history() in cnn_cifar10.py.
• Add the comparison function immediately below plot_history() by pasting this code:

def plot_comparison(histories):
    figure, axes = plt.subplots(1, 2, figsize=(12, 4))
    for variant, history in histories.items():
        epochs = range(1, len(history["train_loss"]) + 1)
        axes[0].plot(
            epochs,
            history["val_loss"],
            marker="o",
            label=f"{variant} validation",
        )
        axes[1].plot(
            epochs,
            history["val_acc"],
            marker="o",
            label=f"{variant} validation",
        )

    axes[0].set_title("Validation loss")
    axes[0].set_xlabel("Epoch")
    axes[0].legend()
    axes[1].set_title("Validation accuracy")
    axes[1].set_xlabel("Epoch")
    axes[1].legend()
    figure.suptitle("Baseline versus improved CNN")
    figure.tight_layout()
    output_path = ARTIFACT_DIR / "model_comparison.png"
    figure.savefig(output_path, dpi=150)
    print(f"Saved {output_path}")
    plt.show()
    plt.close(figure)

What does the comparison show?

• The left chart compares validation loss over the same three epochs.
• The right chart compares validation accuracy over those epochs.
• The function saves the combined result as artifacts/model_comparison.png before displaying it.

• Locate the existing parse_args() function.
• Replace the full function with this completed command-line configuration:

def parse_args():
    parser = argparse.ArgumentParser(
        description="Train and compare compact CIFAR-10 CNNs on a Windows CPU."
    )
    parser.add_argument(
        "--stage",
        choices=("data", "baseline", "compare", "experiment"),
        default="compare",
    )
    parser.add_argument(
        "--variant", choices=("augment", "dropout"), default="augment"
    )
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--train-size", type=int, default=8000)
    parser.add_argument("--val-size", type=int, default=2000)
    parser.add_argument("--test-size", type=int, default=2000)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--retrain", action="store_true")
    return parser.parse_args()

What changes in the command-line interface?

The program now accepts the compare and experiment stages. The size and epoch defaults preserve the same controlled CPU experiment used by the baseline.

• In main(), locate the complete branch that begins with if args.stage == "baseline":.
• Add the experiment branch immediately below that baseline branch:

    if args.stage == "experiment":
        model, history, test_loader = load_or_train(args.variant, args)
        test_loss, test_accuracy = test_model(model, test_loader)
        best_validation = max(history["val_acc"])
        print(
            f"{args.variant} | best val acc {best_validation:.3f} | "
            f"test loss {test_loss:.4f} | test acc {test_accuracy:.3f}"
        )
        return

Why include the experiment branch?

This branch runs one selected regularization variant through the same loader and evaluation path. It reports the best validation accuracy separately from the final test measurements.

• Add the first comparison block immediately below the experiment branch:

    baseline_model, baseline_history, baseline_test_loader = load_or_train(
        "baseline", args
    )
    improved_model, improved_history, improved_test_loader = load_or_train(
        "improved", args
    )

    baseline_test = test_model(baseline_model, baseline_test_loader)
    improved_test = test_model(improved_model, improved_test_loader)
    print(
        f"baseline | test loss {baseline_test[0]:.4f} | "
        f"test acc {baseline_test[1]:.3f}"
    )
    print(
        f"improved | test loss {improved_test[0]:.4f} | "
        f"test acc {improved_test[1]:.3f}"
    )

    histories = {
        "baseline": baseline_history,
        "improved": improved_history,
    }
    plot_comparison(histories)

How does the comparison stay controlled?

Both variants pass through load_or_train() with the same command-line settings. A matching baseline checkpoint is reused. The improved variant trains only when its matching checkpoint is unavailable.

Test loss and accuracy are reported for both models. The histories sent to plot_comparison() contain the validation measurements used for diagnosis.

• Add the model-selection block immediately below plot_comparison(histories).

    baseline_best_val = max(baseline_history["val_acc"])
    improved_best_val = max(improved_history["val_acc"])
    if improved_best_val >= baseline_best_val:
        selected_name = "improved"
        selected_model = improved_model
        selected_loader = improved_test_loader
    else:
        selected_name = "baseline"
        selected_model = baseline_model
        selected_loader = baseline_test_loader

    print(f"Selected by validation accuracy: {selected_name}")
    show_predictions(
        selected_model,
        selected_loader,
        f"Predictions from validation-selected model: {selected_name}",
        "best_model_predictions.png",
    )

Why select with validation accuracy?

The highest value in each val_acc history decides which model supplies the prediction grid. The test results remain final measurements instead of becoming a tuning signal.

Having trouble placing the comparison code?

Keep every new stage branch inside main() with four spaces of indentation. Place the selection block after plot_comparison(histories) so both histories already exist.

Ask for help checking the completed dispatch logic in cnn_cifar10.py.

Run and inspect the comparison

The comparison run may reuse artifacts/baseline.pth because its saved configuration still matches. The improved model trains for three CPU epochs while printing progress after each epoch.

• Save cnn_cifar10.py with Ctrl+S.
• Switch back to the PowerShell terminal from earlier.
• Predict which model will achieve the higher best validation accuracy.
• Run the completed comparison with this command:

python cnn_cifar10.py --stage compare

What happens during this run?

• The matching baseline checkpoint is loaded instead of retrained.
• The improved variant trains with random crops, horizontal flips, and dropout.
• The terminal prints separate baseline and improved test measurements.
• The first plot window compares validation loss and validation accuracy.

• Inspect the first plot window for the relative validation loss and validation accuracy curves.
• Close the comparison plot window to continue the program.
• Read the terminal line beginning with Selected by validation accuracy:.
• Inspect the second plot window for predicted and actual labels from the selected model.
• Close the prediction plot window to finish the run.

You should see artifacts/improved.pth, artifacts/model_comparison.png, and artifacts/best_model_predictions.png in the artifacts folder.

That completes the controlled comparison. Your program now measures whether regularization improved validation accuracy, validation loss, or the gap between training and validation performance.

Comparison run stopped or a plot is missing?

Check the first unfinished epoch or traceback line in PowerShell. A transform-related failure usually points to the new augmented transform. A name-related failure usually points to one of the new comparison blocks.

If the terminal appears paused after saving the comparison chart, close the open plot window. The prediction grid runs after that window closes.

Ask for help diagnosing the comparison run in cnn_cifar10.py.

✔️ Awesome, I've got everything!

Excellent work. Your baseline and improved models now share one controlled comparison path. Confirm that cnn_cifar10.py is saved before continuing.

ⓧ I'd like to double check the full code

Compare your saved cnn_cifar10.py with this complete version.

import argparse
import copy
from pathlib import Path

import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets
from torchvision.transforms import v2

SEED = 42
DEVICE = torch.device("cpu")
DATA_DIR = Path("data")
ARTIFACT_DIR = Path("artifacts")
CLASSES = (
    "airplane",
    "automobile",
    "bird",
    "cat",
    "deer",
    "dog",
    "frog",
    "horse",
    "ship",
    "truck",
)

EVAL_TRANSFORM = v2.Compose(
    [
        v2.ToImage(),
        v2.ToDtype(torch.float32, scale=True),
        v2.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
)

AUGMENTED_TRANSFORM = v2.Compose(
    [
        v2.ToImage(),
        v2.RandomCrop(32, padding=4),
        v2.RandomHorizontalFlip(p=0.5),
        v2.ToDtype(torch.float32, scale=True),
        v2.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
)


class SmallCNN(nn.Module):
    def __init__(self, dropout_probability):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Dropout(p=dropout_probability),
            nn.Linear(128 * 4 * 4, 128),
            nn.ReLU(),
            nn.Dropout(p=dropout_probability),
            nn.Linear(128, len(CLASSES)),
        )

    def forward(self, images):
        return self.classifier(self.features(images))


def transform_for(variant):
    if variant in {"augment", "improved"}:
        return AUGMENTED_TRANSFORM
    return EVAL_TRANSFORM


def dropout_for(variant):
    if variant in {"dropout", "improved"}:
        return 0.4
    return 0.0


def make_loaders(variant, train_size, val_size, test_size, batch_size):
    training_dataset = datasets.CIFAR10(
        root=DATA_DIR,
        train=True,
        download=True,
        transform=transform_for(variant),
    )
    validation_dataset = datasets.CIFAR10(
        root=DATA_DIR,
        train=True,
        download=True,
        transform=EVAL_TRANSFORM,
    )
    test_dataset = datasets.CIFAR10(
        root=DATA_DIR,
        train=False,
        download=True,
        transform=EVAL_TRANSFORM,
    )

    if train_size + val_size > len(training_dataset):
        raise ValueError("train-size plus val-size exceeds the training dataset")

    generator = torch.Generator().manual_seed(SEED)
    shuffled_indices = torch.randperm(
        len(training_dataset), generator=generator
    ).tolist()
    train_indices = shuffled_indices[:train_size]
    val_indices = shuffled_indices[train_size : train_size + val_size]
    test_indices = list(range(min(test_size, len(test_dataset))))

    train_loader = DataLoader(
        Subset(training_dataset, train_indices),
        batch_size=batch_size,
        shuffle=True,
        num_workers=0,
    )
    val_loader = DataLoader(
        Subset(validation_dataset, val_indices),
        batch_size=batch_size,
        shuffle=False,
        num_workers=0,
    )
    test_loader = DataLoader(
        Subset(test_dataset, test_indices),
        batch_size=batch_size,
        shuffle=False,
        num_workers=0,
    )
    return train_loader, val_loader, test_loader


def run_pass(model, loader, loss_function, optimizer=None):
    training = optimizer is not None
    if training:
        model.train()
    else:
        model.eval()

    total_loss = 0.0
    correct = 0
    total = 0

    with torch.set_grad_enabled(training):
        for images, labels in loader:
            images = images.to(DEVICE)
            labels = labels.to(DEVICE)

            if training:
                optimizer.zero_grad(set_to_none=True)

            logits = model(images)
            loss = loss_function(logits, labels)

            if training:
                loss.backward()
                optimizer.step()

            predictions = logits.argmax(dim=1)
            total_loss += loss.item() * labels.size(0)
            correct += (predictions == labels).sum().item()
            total += labels.size(0)

    return total_loss / total, correct / total


def train_variant(variant, train_loader, val_loader, epochs):
    torch.manual_seed(SEED)
    model = SmallCNN(dropout_for(variant)).to(DEVICE)
    loss_function = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    history = {"train_loss": [], "val_loss": [], "train_acc": [], "val_acc": []}
    best_validation_accuracy = -1.0
    best_state = None

    for epoch in range(1, epochs + 1):
        train_loss, train_accuracy = run_pass(
            model, train_loader, loss_function, optimizer
        )
        val_loss, val_accuracy = run_pass(model, val_loader, loss_function)

        history["train_loss"].append(train_loss)
        history["val_loss"].append(val_loss)
        history["train_acc"].append(train_accuracy)
        history["val_acc"].append(val_accuracy)

        print(
            f"{variant:>9} | epoch {epoch}/{epochs} | "
            f"train loss {train_loss:.4f} | val loss {val_loss:.4f} | "
            f"train acc {train_accuracy:.3f} | val acc {val_accuracy:.3f}"
        )

        if val_accuracy > best_validation_accuracy:
            best_validation_accuracy = val_accuracy
            best_state = copy.deepcopy(model.state_dict())

    model.load_state_dict(best_state)
    return model, history


def load_or_train(variant, args):
    train_loader, val_loader, test_loader = make_loaders(
        variant,
        args.train_size,
        args.val_size,
        args.test_size,
        args.batch_size,
    )
    checkpoint_path = ARTIFACT_DIR / f"{variant}.pth"
    expected_config = {
        "variant": variant,
        "epochs": args.epochs,
        "train_size": args.train_size,
        "val_size": args.val_size,
    }

    if checkpoint_path.exists() and not args.retrain:
        checkpoint = torch.load(
            checkpoint_path, map_location=DEVICE, weights_only=True
        )
        if checkpoint.get("config") == expected_config:
            model = SmallCNN(dropout_for(variant)).to(DEVICE)
            model.load_state_dict(checkpoint["model_state_dict"])
            print(f"Reused {checkpoint_path}")
            return model, checkpoint["history"], test_loader

    model, history = train_variant(variant, train_loader, val_loader, args.epochs)
    torch.save(
        {
            "model_state_dict": model.state_dict(),
            "history": history,
            "config": expected_config,
        },
        checkpoint_path,
    )
    print(f"Saved {checkpoint_path}")
    return model, history, test_loader


def test_model(model, test_loader):
    loss_function = nn.CrossEntropyLoss()
    return run_pass(model, test_loader, loss_function)


def display_image(image):
    return (image * 0.5 + 0.5).clamp(0, 1).permute(1, 2, 0).numpy()


def show_sample_grid():
    dataset = datasets.CIFAR10(
        root=DATA_DIR,
        train=True,
        download=True,
        transform=EVAL_TRANSFORM,
    )
    examples = {}
    for image, label in dataset:
        if label not in examples:
            examples[label] = image
        if len(examples) == len(CLASSES):
            break

    figure, axes = plt.subplots(2, 5, figsize=(11, 5))
    for label, axis in enumerate(axes.flat):
        axis.imshow(display_image(examples[label]))
        axis.set_title(CLASSES[label])
        axis.axis("off")
    figure.suptitle("CIFAR-10 sample classes")
    figure.tight_layout()
    output_path = ARTIFACT_DIR / "cifar10_samples.png"
    figure.savefig(output_path, dpi=150)
    print(f"Saved {output_path}")
    plt.show()
    plt.close(figure)


def plot_history(history, title, output_name):
    epochs = range(1, len(history["train_loss"]) + 1)
    figure, axes = plt.subplots(1, 2, figsize=(11, 4))
    axes[0].plot(epochs, history["train_loss"], marker="o", label="train")
    axes[0].plot(epochs, history["val_loss"], marker="o", label="validation")
    axes[0].set_title("Loss")
    axes[0].set_xlabel("Epoch")
    axes[0].legend()

    axes[1].plot(epochs, history["train_acc"], marker="o", label="train")
    axes[1].plot(epochs, history["val_acc"], marker="o", label="validation")
    axes[1].set_title("Accuracy")
    axes[1].set_xlabel("Epoch")
    axes[1].legend()

    figure.suptitle(title)
    figure.tight_layout()
    output_path = ARTIFACT_DIR / output_name
    figure.savefig(output_path, dpi=150)
    print(f"Saved {output_path}")
    plt.show()
    plt.close(figure)


def plot_comparison(histories):
    figure, axes = plt.subplots(1, 2, figsize=(12, 4))
    for variant, history in histories.items():
        epochs = range(1, len(history["train_loss"]) + 1)
        axes[0].plot(
            epochs,
            history["val_loss"],
            marker="o",
            label=f"{variant} validation",
        )
        axes[1].plot(
            epochs,
            history["val_acc"],
            marker="o",
            label=f"{variant} validation",
        )

    axes[0].set_title("Validation loss")
    axes[0].set_xlabel("Epoch")
    axes[0].legend()
    axes[1].set_title("Validation accuracy")
    axes[1].set_xlabel("Epoch")
    axes[1].legend()
    figure.suptitle("Baseline versus improved CNN")
    figure.tight_layout()
    output_path = ARTIFACT_DIR / "model_comparison.png"
    figure.savefig(output_path, dpi=150)
    print(f"Saved {output_path}")
    plt.show()
    plt.close(figure)


def show_predictions(model, test_loader, title, output_name):
    model.eval()
    images, labels = next(iter(test_loader))
    with torch.no_grad():
        predictions = model(images.to(DEVICE)).argmax(dim=1).cpu()

    figure, axes = plt.subplots(2, 4, figsize=(10, 5))
    for index, axis in enumerate(axes.flat):
        predicted_name = CLASSES[predictions[index].item()]
        actual_name = CLASSES[labels[index].item()]
        axis.imshow(display_image(images[index]))
        axis.set_title(f"P: {predicted_name}\nA: {actual_name}")
        axis.axis("off")

    figure.suptitle(title)
    figure.tight_layout()
    output_path = ARTIFACT_DIR / output_name
    figure.savefig(output_path, dpi=150)
    print(f"Saved {output_path}")
    plt.show()
    plt.close(figure)


def parse_args():
    parser = argparse.ArgumentParser(
        description="Train and compare compact CIFAR-10 CNNs on a Windows CPU."
    )
    parser.add_argument(
        "--stage",
        choices=("data", "baseline", "compare", "experiment"),
        default="compare",
    )
    parser.add_argument(
        "--variant", choices=("augment", "dropout"), default="augment"
    )
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--train-size", type=int, default=8000)
    parser.add_argument("--val-size", type=int, default=2000)
    parser.add_argument("--test-size", type=int, default=2000)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--retrain", action="store_true")
    return parser.parse_args()


def main():
    args = parse_args()
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    print(f"Device: {DEVICE}")

    if args.stage == "data":
        show_sample_grid()
        return

    if args.stage == "baseline":
        model, history, test_loader = load_or_train("baseline", args)
        test_loss, test_accuracy = test_model(model, test_loader)
        print(f"baseline | test loss {test_loss:.4f} | test acc {test_accuracy:.3f}")
        plot_history(history, "Baseline CNN", "baseline_curves.png")
        show_predictions(
            model,
            test_loader,
            "Baseline CNN predictions",
            "baseline_predictions.png",
        )
        return

    if args.stage == "experiment":
        model, history, test_loader = load_or_train(args.variant, args)
        test_loss, test_accuracy = test_model(model, test_loader)
        best_validation = max(history["val_acc"])
        print(
            f"{args.variant} | best val acc {best_validation:.3f} | "
            f"test loss {test_loss:.4f} | test acc {test_accuracy:.3f}"
        )
        return

    baseline_model, baseline_history, baseline_test_loader = load_or_train(
        "baseline", args
    )
    improved_model, improved_history, improved_test_loader = load_or_train(
        "improved", args
    )

    baseline_test = test_model(baseline_model, baseline_test_loader)
    improved_test = test_model(improved_model, improved_test_loader)
    print(
        f"baseline | test loss {baseline_test[0]:.4f} | "
        f"test acc {baseline_test[1]:.3f}"
    )
    print(
        f"improved | test loss {improved_test[0]:.4f} | "
        f"test acc {improved_test[1]:.3f}"
    )

    histories = {
        "baseline": baseline_history,
        "improved": improved_history,
    }
    plot_comparison(histories)

    baseline_best_val = max(baseline_history["val_acc"])
    improved_best_val = max(improved_history["val_acc"])
    if improved_best_val >= baseline_best_val:
        selected_name = "improved"
        selected_model = improved_model
        selected_loader = improved_test_loader
    else:
        selected_name = "baseline"
        selected_model = baseline_model
        selected_loader = baseline_test_loader

    print(f"Selected by validation accuracy: {selected_name}")
    show_predictions(
        selected_model,
        selected_loader,
        f"Predictions from validation-selected model: {selected_name}",
        "best_model_predictions.png",
    )


if __name__ == "__main__":
    main()

Secret mission

Run an Ablation Study

The improved CNN combines augmentation plus dropout, which hides each technique's individual contribution. Run two controlled experiments to discover which change helps most on your machine.

Clean Up Your Resources

Clean Up Your Resources

Everything you created stays on your Windows laptop with no ongoing cost. Decide whether to keep the project available, pause your work, or delete the local files.

Resources you used:

• The cnn-cifar10-local project folder containing cnn_cifar10.py.
• The cnn-cifar10-local/.venv virtual environment containing PyTorch 2.14.0, TorchVision 0.29.0, and Matplotlib 3.11.1.
• The data folder containing the downloaded CIFAR-10 training and test data.
• The artifacts folder containing sample images, learning curves, comparison plots, prediction grids, and .pth checkpoints.

Keep everything running

No action needed. Choose this if you are still comparing the four model variants or want to rerun an experiment.

• Leave Visual Studio Code open with cnn-cifar10-local/.venv selected.
• Keep cnn-cifar10-local in place so later runs can reuse data and artifacts.

Pause - I'll come back to this later

Shut down your local workspace to clear the screen. Your project files stay unchanged for the next session.

• Close any Matplotlib plot windows.
• Close Visual Studio Code.

Delete - I don't want to use this again

Remove all local project resources when you want a clean reset. Your charts and checkpoints stay safe if you copy them outside the project folder first.

• Copy any charts or checkpoints you want to keep from artifacts to a folder outside cnn-cifar10-local.
• Press the Windows key to open search.
• Type File Explorer.
• Press Enter to open File Explorer.
• Locate the cnn-cifar10-local folder.
• Delete the cnn-cifar10-local folder through File Explorer.
• Confirm that cnn-cifar10-local no longer appears in File Explorer.

Deleting cnn-cifar10-local removes cnn_cifar10.py, cnn-cifar10-local/.venv, data, and artifacts in one action.

Nice Work!

Nice Work!

You did it! Your local PyTorch program now classifies CIFAR-10 images with a model selected by validation accuracy.

You've learned how to:

• Build a compact convolutional neural network for 10-category CIFAR-10 classification on a Windows CPU.
• Diagnose generalization with separate training metrics plus validation metrics. Read saved learning curves to spot whether stronger training performance carries over to unseen images.
• Test data augmentation plus dropout against a controlled baseline. Keep reusable model checkpoints plus charts in artifacts for future experiments.
• Secret Mission: Isolated data augmentation from dropout in a controlled ablation study. Recorded the most useful change with a caveat about the CPU-sized subset.

Ready to quiz yourself?