CNN Image Classification — Animal Faces (PyTorch)

A compact convolutional neural network that classifies Animal Faces (AFHQ) images into three categories — cat, dog, wild — implemented in PyTorch with a clean, reproducible training pipeline (data augmentation, normalization, seeded DataLoaders, train/validation evaluation).

Dataset

The Animal Faces (AFHQ) dataset: ~16,130 32×32 RGB images across 3 balanced classes (cat / dog / wild), in torchvision ImageFolder layout:

data/afhq/train/<class>/*.png
data/afhq/val/<class>/*.png

data.py can download it automatically (--download) or use a local copy. Training images are augmented with random rotation (±10°) and horizontal flips; all images are normalized to the range [−1, 1].

Model

A small CNN (model.py) — three convolutional blocks growing channel depth while halving spatial resolution, then global average pooling and a linear head:

input (3, 32, 32)
  -> [Conv3x3 -> BN -> ReLU -> Conv3x3 -> BN -> ReLU -> MaxPool] x3   (32 -> 64 -> 128 channels)
  -> global average pool -> Dropout -> Linear -> 3 logits

~288K parameters. Batch normalization stabilizes training; global average pooling keeps the head small and makes the network input-size agnostic.

Usage

pip install -r requirements.txt

# download AFHQ and train
python train.py --download --epochs 20 --save cnn.pt

# or point at a local copy of the dataset
python train.py --data-dir ./data --epochs 20

# quick end-to-end smoke test with synthetic images (no download)
python train.py --fake --epochs 1

The script prints train/validation accuracy each epoch. Chance accuracy for this 3-class task is 33.3%.

Repository contents

File	Description
model.py	Compact CNN classifier (conv blocks + global average pooling)
data.py	AFHQ download + augmented / normalized DataLoaders
train.py	Training / evaluation loop with CLI arguments (incl. --fake smoke test)

Tech stack

Python · PyTorch · torchvision · NumPy

Acknowledgment

Built as part of coursework for Artificial Neural Networks for Neuroscience. The Animal Faces (AFHQ) dataset and data pipeline follow the course materials; the CNN classifier here is an original implementation.