PyTorch CNN Image Classifier
A modular, production-ready Convolutional Neural Network (CNN) architecture implemented in PyTorch for multi-class image classification.
Features
- Modular PyTorch Architecture (
model.py): 3 convolutional blocks with Batch Normalization, ReLU, MaxPool, Dropout, Adaptive Average Pooling, and a classification head. - Ready-to-Use Inference Pipeline (
predict.py): Predict directly on images or run instant synthetic demo tests with top-k ranked probabilities and CIFAR-10 labels. - Automated Test Suite (
test_pipeline.py): Unit tests validating forward pass, variable resolution handling, gradient backpropagation flow, and checkpoint saving/loading. - Interactive Local Web App (
app.py): Built-in Gradio interface to test and classify uploaded images in your browser. - Config & Label Mapping (
config.json): Formal model metadata withid2labelmapping for the 10 classes. - Training Script (
train.py): Full training and validation pipeline on CIFAR-10 with data augmentation, Cosine Annealing learning rate scheduling, and automatic best-model checkpointing.
Quickstart
1. Installation
git clone https://huggingface.co/sarvesh74/pytorch-cnn-image-classifier
cd pytorch-cnn-image-classifier
pip install -r requirements.txt
2. Verify Installation & Architecture
Run the test suite to verify that all layers, tensor dimensions, and backward passes function as expected:
python test_pipeline.py
3. Run Inference
Instant Synthetic Demo
Test inference immediately without needing an external image file:
python predict.py --demo
Classify an Image File
python predict.py --image path/to/your/image.jpg --top-k 3
Python API Usage
from PIL import Image
from predict import ImageClassifier
classifier = ImageClassifier()
image = Image.open("your_image.jpg")
predictions = classifier.predict(image, top_k=3)
for pred in predictions:
print(f"{pred['label']}: {pred['confidence']:.2f}%")
4. Launch the Interactive Web App (Gradio)
To launch an interactive browser UI where you can upload and test images:
python app.py
Open the provided local URL (typically http://127.0.0.1:7860) in your browser.
5. Training on CIFAR-10
Train the network from scratch using the included pipeline:
python train.py --epochs 20 --batch-size 64 --lr 0.001
The script will automatically download the dataset, track validation accuracy across epochs, and export the best checkpoint to best_model.pth.
Model Architecture Details
| Layer Stage | Input Size | Details | Output Channels |
|---|---|---|---|
| Conv Block 1 | $3 \times W \times H$ | Conv2d ($3\times3$, $P=1$) $\to$ BN $\to$ ReLU $\to$ Conv2d $\to$ BN $\to$ ReLU $\to$ MaxPool ($2\times2$) $\to$ Dropout | 32 |
| Conv Block 2 | $32 \times \frac{W}{2} \times \frac{H}{2}$ | Conv2d ($3\times3$, $P=1$) $\to$ BN $\to$ ReLU $\to$ Conv2d $\to$ BN $\to$ ReLU $\to$ MaxPool ($2\times2$) $\to$ Dropout | 64 |
| Conv Block 3 | $64 \times \frac{W}{4} \times \frac{H}{4}$ | Conv2d ($3\times3$, $P=1$) $\to$ BN $\to$ ReLU $\to$ MaxPool ($2\times2$) $\to$ Dropout | 128 |
| Pooling | $128 \times \frac{W}{8} \times \frac{H}{8}$ | AdaptiveAvgPool2d((4, 4)) |
$128 \times 4 \times 4$ |
| Classifier | 2048 | Linear(2048, 512) $\to$ BN $\to$ ReLU $\to$ Dropout(0.5) $\to$ Linear(512, 10) | 10 classes |
Spatial Dimension Formula
- Downloads last month
- 43