This project implements a lightweight CNN model for CIFAR-10 image classification using PyTorch. The architecture is designed to achieve high accuracy with minimal parameters (~162k) by combining:
- Depthwise Separable Convolutions: Reduces redundant computations, inspired by MobileNet.
- Dilated Convolutions: Expands the receptive field without pooling or increasing parameters.
- Adaptive Average Pooling and Fully Connected Layer: For classification.
The goal is to demonstrate a balance between efficiency and performance on a standard benchmark dataset.
CIFAR-10 consists of 60,000 32x32 color images across 10 classes:
-
airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck
-
Train: 50,000 images
-
Test: 10,000 images
Data is loaded using torchvision.datasets.CIFAR10 with Albumentations augmentations.
Data augmentations improve generalization and reduce overfitting.
| Augmentation | Description | Probability |
|---|---|---|
| Horizontal Flip | Flips the image along vertical axis | 0.5 |
| Shift, Scale, Rotate | Small translations, scaling & rotations | 0.5 |
| Coarse Dropout | Randomly masks 16x16 patches | 1 |
| Normalize | Standardize pixel values to dataset mean & std | 1 |
Train transforms: HorizontalFlip, ShiftScaleRotate, CoarseDropout, Normalize, ToTensor
Test transforms: Normalize, ToTensor
The model, CustomNet, consists of 4 convolutional blocks followed by Global Average Pooling (GAP) and a Linear layer.
Block-wise Description:
| Block | Layers | Purpose | Output Size |
|---|---|---|---|
| C1 | 2 Conv + BatchNorm + ReLU | Extract low-level features (edges/textures) | 28x28x24 |
| C2 | Depthwise Separable Conv + Conv + BN + ReLU | Reduce spatial dimension, increase channels | 14x14x40 |
| C3 | Dilated Conv + Conv + BN + ReLU | Capture larger context | 7x7x80 |
| C4 | Conv + BN + ReLU + Conv | Prepare for classification | 7x7x64 |
| GAP + FC | AdaptiveAvgPool + Linear | Convert feature map to class logits | 1x1x10 |
| Layer | Kernel | Stride | Padding | Dilation | RF | Jump |
|---|---|---|---|---|---|---|
| C1-1 | 3 | 1 | 0 | 1 | 3 | 1 |
| C1-2 | 3 | 1 | 0 | 1 | 5 | 1 |
| C2-1 (Depthwise) | 3 | 1 | 2 | 2 | 9 | 1 |
| C2-2 | 3 | 2 | 1 | 1 | 11 | 2 |
| C3-1 | 3 | 1 | 4 | 4 | 27 | 2 |
| C3-2 | 3 | 2 | 1 | 1 | 31 | 4 |
| C4-1 | 3 | 1 | 3 | 3 | 55 | 4 |
| C4-2 | 1 | 1 | 0 | 1 | 55 | 4 |
| Layer | H | W |
|---|---|---|
| C1-1 | 30 | 30 |
| C1-2 | 28 | 28 |
| C2-1 (Depthwise) | 28 | 28 |
| C2-2 | 14 | 14 |
| C3-1 | 14 | 14 |
| C3-2 | 7 | 7 |
| C4-1 | 7 | 7 |
| C4-2 | 7 | 7 |
| Layer (type) | Output Shape | Param # |
|---|---|---|
| Conv2d | [1, 16, 30, 30] | 448 |
| BatchNorm2d | [1, 16, 30, 30] | 32 |
| ReLU | [1, 16, 30, 30] | -- |
| Conv2d | [1, 24, 28, 28] | 3,480 |
| BatchNorm2d | [1, 24, 28, 28] | 48 |
| ReLU | [1, 24, 28, 28] | -- |
| DepthwiseSeparableConv | [1, 32, 28, 28] | 1,040 |
| BatchNorm2d | [1, 32, 28, 28] | 64 |
| ReLU | [1, 32, 28, 28] | -- |
| Conv2d | [1, 40, 14, 14] | 11,560 |
| BatchNorm2d | [1, 40, 14, 14] | 80 |
| ReLU | [1, 40, 14, 14] | -- |
| Conv2d | [1, 64, 14, 14] | 23,104 |
| BatchNorm2d | [1, 64, 14, 14] | 128 |
| ReLU | [1, 64, 14, 14] | -- |
| Conv2d | [1, 80, 7, 7] | 46,160 |
| BatchNorm2d | [1, 80, 7, 7] | 160 |
| ReLU | [1, 80, 7, 7] | -- |
| Conv2d | [1, 96, 7, 7] | 69,216 |
| BatchNorm2d | [1, 96, 7, 7] | 192 |
| ReLU | [1, 96, 7, 7] | -- |
| Conv2d | [1, 64, 7, 7] | 6,208 |
| AdaptiveAvgPool2d | [1, 64, 1, 1] | -- |
| Linear | [1, 10] | 650 |
Total params: 162,570
Trainable params: 162,570
Non-trainable params: 0
Estimated Total Size (MB): 2.24
- Device: GPU if available else CPU
- Optimizer: SGD (lr=0.05, momentum=0.9, weight decay=5e-4)
- Scheduler: OneCycleLR (
max_lr=0.05,pct_start=0.2, cosine annealing) - Epochs: 30
- Batch Size: 128
- Loss Function: CrossEntropyLoss
Training Loop Highlights:
- Tracks train and test loss/accuracy per epoch
- Prints accuracy & loss difference to monitor overfitting
- Saves the best test accuracy
| Epoch | Train Loss | Train Acc (%) | Test Loss | Test Acc (%) | Acc Diff | Loss Diff |
|---|---|---|---|---|---|---|
| 1 | 1.7043 | 36.1 | 1.3578 | 49.78 | -13.68 | 0.3465 |
| 2 | 1.2413 | 55.1 | 1.2694 | 55.02 | 0.08 | -0.0281 |
| 3 | 1.0397 | 62.72 | 0.9515 | 66 | -3.28 | 0.0882 |
| 4 | 0.9106 | 67.68 | 0.9627 | 67.14 | 0.54 | -0.0521 |
| 5 | 0.8304 | 70.66 | 0.906 | 68.8 | 1.86 | -0.0756 |
| 6 | 0.7705 | 72.93 | 0.8015 | 71.96 | 0.97 | -0.031 |
| 7 | 0.7262 | 74.56 | 0.833 | 71.79 | 2.77 | -0.1068 |
| 8 | 0.6948 | 75.87 | 0.7104 | 75.59 | 0.28 | -0.0155 |
| 9 | 0.6628 | 76.95 | 0.6434 | 77.78 | -0.83 | 0.0194 |
| 10 | 0.6426 | 77.53 | 0.6097 | 78.83 | -1.3 | 0.0329 |
| 11 | 0.6226 | 78.34 | 0.6237 | 78.3 | 0.04 | -0.0011 |
| 12 | 0.6065 | 78.9 | 0.6541 | 77.45 | 1.45 | -0.0476 |
| 13 | 0.5922 | 79.35 | 0.7029 | 75.8 | 3.55 | -0.1107 |
| 14 | 0.5753 | 79.8 | 0.6217 | 78.79 | 1.01 | -0.0464 |
| 15 | 0.5635 | 80.32 | 0.7094 | 75.56 | 4.76 | -0.146 |
| 16 | 0.5541 | 80.46 | 0.6229 | 78.94 | 1.52 | -0.0689 |
| 17 | 0.5396 | 81.02 | 0.5751 | 79.89 | 1.13 | -0.0355 |
| 18 | 0.5224 | 81.78 | 0.5399 | 81.45 | 0.33 | -0.0175 |
| 19 | 0.515 | 82.03 | 0.564 | 80.55 | 1.48 | -0.0489 |
| 20 | 0.5014 | 82.53 | 0.5168 | 82.24 | 0.29 | -0.0154 |
| 21 | 0.4848 | 82.97 | 0.5259 | 82.08 | 0.89 | -0.0411 |
| 22 | 0.4616 | 84.01 | 0.5039 | 82.62 | 1.39 | -0.0424 |
| 23 | 0.4485 | 84.32 | 0.501 | 83.32 | 1 | -0.0524 |
| 24 | 0.4272 | 85.08 | 0.4734 | 83.32 | 1.76 | -0.0462 |
| 25 | 0.4062 | 85.87 | 0.4669 | 84.03 | 1.84 | -0.0607 |
| 26 | 0.382 | 86.7 | 0.4561 | 84.45 | 2.25 | -0.0741 |
| 27 | 0.3631 | 87.22 | 0.4357 | 85.29 | 1.93 | -0.0726 |
| 28 | 0.3432 | 88.12 | 0.421 | 85.56 | 2.56 | -0.0778 |
| 29 | 0.3281 | 88.67 | 0.4207 | 85.62 | 3.05 | -0.0925 |
| 30 | 0.3229 | 88.96 | 0.4186 | 85.72 | 3.24 | -0.0957 |
- Clone the repository:
git clone <repo-url>
cd <repo-folder>- Set up a virtual environment
python -m venv venv
source venv/bin/activate # On Windows: venv\Scripts\activate- Install dependencies
pip install -r requirements.txt- Run the training script
python cifar.py