🔥 PyTorch Image Classifier

A modular, configurable image classification framework built with PyTorch. Supports custom CNNs and transfer learning from pretrained models (ResNet, EfficientNet, MobileNet).

---

✨ Features

• 🧱 Modular architecture — swap models, datasets, and optimizers via config
• 🔁 Transfer learning — fine-tune ResNet18/50, EfficientNet-B0, MobileNetV3
• 📊 Training dashboard — live loss/accuracy curves saved as plots
• 🧪 Evaluation suite — confusion matrix, per-class accuracy, top-k metrics
• ⚙️ YAML config — no code changes needed to run new experiments
• 💾 Checkpointing — auto-saves best model, resume training anytime
• 🪵 Logging — structured logs with Python logging + optional TensorBoard

---

📁 Project Structure

pytorch-image-classifier/
├── configs/
│   └── default.yaml          # Default experiment config
├── src/
│   ├── dataset.py            # Dataset loading & augmentation
│   ├── model.py              # Model factory (custom CNN + pretrained)
│   ├── trainer.py            # Training & validation loop
│   ├── evaluator.py          # Metrics, confusion matrix, plots
│   └── utils.py              # Helpers, logging, checkpointing
├── scripts/
│   ├── train.py              # Entry point for training
│   └── predict.py            # Run inference on new images
├── tests/
│   └── test_model.py         # Unit tests
├── requirements.txt
└── README.md

---

🚀 Quick Start

1. Install dependencies

pip install -r requirements.txt

2. Prepare your dataset

Organize images in this structure:

data/
├── train/
│   ├── cat/
│   └── dog/
└── val/
    ├── cat/
    └── dog/

💡 Tip: Works out of the box with CIFAR-10 or any ImageFolder-compatible dataset.

3. Configure your experiment

Edit configs/default.yaml or create a new config file:

model:
  name: resnet18       # Options: custom_cnn, resnet18, resnet50, efficientnet_b0, mobilenet_v3
  pretrained: true
  num_classes: 10

training:
  epochs: 30
  batch_size: 64
  learning_rate: 0.001
  optimizer: adam      # Options: adam, sgd, adamw

4. Train

python scripts/train.py --config configs/default.yaml

5. Evaluate / Predict

python scripts/predict.py --config configs/default.yaml --image path/to/image.jpg

---

📈 Results

Training curves and confusion matrices are automatically saved to outputs/.

---

🔧 Supported Models

Model	Pretrained	Notes
custom_cnn	❌	Lightweight 4-layer CNN
resnet18	✅	Fast, great for small datasets
resnet50	✅	Higher accuracy, more compute
efficientnet_b0	✅	Efficient & accurate
mobilenet_v3	✅	Mobile-friendly

---

🧪 Run Tests

pytest tests/

---

📄 License

MIT License — free to use, modify, and distribute.