🖼️ CIFAR-10 Image Classifier — CNN vs YOLOv8

Image classification on CIFAR-10 using a custom PyTorch CNN and YOLOv8 — a side-by-side comparison of a handcrafted deep learning pipeline vs a state-of-the-art pretrained model.

---

🚀 Project Overview

This project tackles image classification on the CIFAR-10 dataset — 60,000 color images across 10 categories — using two distinct deep learning approaches:

Approach	Model	Strategy
🔨 Custom CNN	SimpleCNN (PyTorch)	Designed and trained from scratch
⚡ YOLO Classifier	YOLOv8n-cls (Ultralytics)	Fine-tuned pretrained backbone

The goal is to explore how a carefully architected custom CNN compares against a modern pretrained model on a classic benchmark — with a full training, evaluation, and prediction pipeline.

10 Classes: airplane · automobile · bird · cat · deer · dog · frog · horse · ship · truck

---

🧠 Model Architecture

SimpleCNN (PyTorch — built from scratch)

Input (3×32×32)
    │
    ├── Conv2d(3 → 128, 3×3) + BatchNorm + ReLU
    ├── Conv2d(128 → 256, 3×3) + BatchNorm + ReLU + MaxPool2d(2×2)
    ├── Dropout2d(0.1)
    ├── Conv2d(256 → 256, 3×3) + BatchNorm + ReLU
    ├── Conv2d(256 → 128, 3×3) + BatchNorm + ReLU + MaxPool2d(2×2)
    │
    └── Flatten → Linear(8192 → 512) → ReLU → Dropout(0.3) → Linear(512 → 10)

Training config:

Parameter	Value
Optimizer	SGD + Nesterov
Learning Rate	0.1
Momentum	0.9
Weight Decay	1e-4
Scheduler	CosineAnnealingLR
Loss	CrossEntropyLoss (label smoothing 0.1)
Epochs	18
Batch Size	128

YOLOv8n-cls (Ultralytics)

Fine-tuned on CIFAR-10 images resized to 64×64, organized into a folder-based dataset structure. Trained for 7 epochs with batch size 64.

---

📊 Results

Model	Train Accuracy	Test Accuracy (Top-1)	Top-5 Accuracy
SimpleCNN	91%	89%	—
YOLOv8n-cls	—	85.46%	99.49%

💡 The custom CNN reaches 91% train / 89% test accuracy, showing strong generalization with minimal overfitting. YOLOv8 achieves 85.46% Top-1 and a near-perfect 99.49% Top-5, meaning the correct class is almost always in the model's top 5 predictions.

On loss: the SimpleCNN training loss decreases steadily across 18 epochs thanks to CosineAnnealingLR scheduling and label smoothing, which prevents overconfident predictions and keeps the loss smooth. Loss and accuracy curves are plotted at the end of training for full transparency.

---

🛠️ Tech Stack

Tool	Role
Python 3.10+	Core language
PyTorch 2.x	Custom CNN training
Torchvision	Dataset loading & transforms
Ultralytics YOLOv8	YOLO classification
Matplotlib	Visualization
Pillow	Image processing
Google Colab	Training environment

---

⚙️ Installation & Setup

1. Clone the repository

git clone https://github.com/your-username/cifar10-cnn-yolo-classifier.git
cd cifar10-cnn-yolo-classifier

2. Install dependencies

pip install torch torchvision matplotlib ultralytics pillow

3. Open in Google Colab

Upload CNN_cifar10.ipynb and run all cells — the CIFAR-10 dataset is downloaded automatically.

---

Predict on custom images

Upload your own images directly in Colab — the model returns the Top-3 predictions with confidence scores.

predict_yolo(yolo_model, class_names)  # YOLOv8 prediction
# or run cell 9 for SimpleCNN prediction

---

📂 Project Structure

cifar10-cnn-yolo-classifier/
│
├── CNN_cifar10.ipynb        # Main notebook (CNN + YOLOv8 pipeline)
├── README.md                # Project documentation

---

🔍 Key Features

• ✅ End-to-end pipeline — from raw data to prediction in a single notebook
• ✅ Custom CNN architecture — designed and trained from scratch with PyTorch
• ✅ YOLOv8 fine-tuning — fast convergence with a pretrained backbone
• ✅ Data augmentation — random crop + horizontal flip during training
• ✅ Learning curves — loss and accuracy plots after training
• ✅ Custom image prediction — upload any image and get instant results
• ✅ GPU-ready — automatic CUDA detection for faster training

---

📈 Performance Insights

• SimpleCNN benefits from BatchNorm at every conv layer — stabilizes training and allows higher learning rates.
• Label smoothing (0.1) reduces overconfidence and improves generalization.
• CosineAnnealingLR provides smooth learning rate decay without manual scheduling.
• YOLOv8 reaches solid accuracy in just 7 epochs thanks to ImageNet-pretrained weights.
• Gradient clipping (max_norm=1.0) prevents exploding gradients during CNN training.

---

🤝 Contribution

Contributions are welcome! Feel free to open an issue or submit a pull request.

# Fork → Clone → Create branch → Commit → Push → Pull Request
git checkout -b feature/your-feature
git commit -m "Add: your feature description"
git push origin feature/your-feature

---

Made with ❤️ and PyTorch