Image Classifier

A small hands-on project for practicing computer vision fundamentals: a convolutional neural network that classifies images from the CIFAR-10 dataset, built with PyTorch.

Overview

CIFAR-10 consists of 60,000 32×32 color images across 10 classes (airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck). This project implements a small CNN with a scripted training pipeline (src/) and an exploratory Jupyter notebook (notebooks/).

Project Structure

.
├── Dockerfile
├── requirements.txt
├── data/                         # CIFAR-10 dataset (downloaded automatically)
├── models/                       # Saved checkpoints (created at train time)
├── notebooks/
│   └── train_nn_cifar_10.ipynb   # End-to-end training walkthrough
└── src/
    ├── main.py                   # CLI entrypoint (train → evaluate → save)
    ├── model.py                  # CNN architecture (Net)
    ├── dataset.py                # DatasetTransformer wrapper around torchvision
    ├── train.py                  # Training loop
    └── evaluate.py               # Evaluation / accuracy

Model

Net (defined in src/model.py) is a simple CNN:

1. Conv2d (3 → 8 channels, 3×3 kernel) + ReLU + MaxPool
2. Conv2d (8 → 16 channels, 3×3 kernel) + ReLU + MaxPool
3. Fully connected layer (16 × 8 × 8 → 10 class scores)

Training uses GPU automatically when CUDA is available; otherwise it falls back to CPU.

Setup

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Usage

Local training

From the src/ directory:

PYTHONPATH=../ python3 main.py \
  --epochs 20 \
  --batch_size 64 \
  --learning_rate 0.001 \
  --data_path ../data \
  --model_path models/cifar_net.pt

Common flags:

Flag	Default	Description
--epochs	20	Number of training epochs
--batch_size	64	Mini-batch size
--learning_rate	0.001	Adam learning rate
--num_workers	4	DataLoader workers
--data_path	../data	CIFAR-10 root directory
--model_path	models/cifar_net.pt	Where to save the checkpoint

Docker

Build:

docker build -t image-classifier .

Run (persists dataset and model on the host):

docker run --rm \
  -v "$(pwd)/data:/app/data" \
  -v "$(pwd)/models:/app/models" \
  image-classifier

Override training args:

docker run --rm \
  -v "$(pwd)/data:/app/data" \
  -v "$(pwd)/models:/app/models" \
  image-classifier \
  python src/main.py --epochs 5 --batch_size 64 \
    --data_path /app/data --model_path /app/models/cifar_net.pt

The default image is CPU-based. For GPU, use a CUDA PyTorch base image (see comments in the Dockerfile) and run with --gpus all.

Notebooks

For a guided walkthrough of loading data, defining the network, training, and evaluating it, see notebooks/train_nn_cifar_10.ipynb.