https://github.com/ccomkhj/lightening_classifier
PyTorch Lightning wrapper to make training classifiers easier.
https://github.com/ccomkhj/lightening_classifier
classification classifier computer-vision ml pytorch pytorch-lightning resnet swin
Last synced: about 2 months ago
JSON representation
PyTorch Lightning wrapper to make training classifiers easier.
- Host: GitHub
- URL: https://github.com/ccomkhj/lightening_classifier
- Owner: ccomkhj
- License: mit
- Created: 2025-02-20T14:17:49.000Z (over 1 year ago)
- Default Branch: main
- Last Pushed: 2025-03-01T14:47:37.000Z (over 1 year ago)
- Last Synced: 2025-09-09T23:34:09.748Z (10 months ago)
- Topics: classification, classifier, computer-vision, ml, pytorch, pytorch-lightning, resnet, swin
- Language: Jupyter Notebook
- Homepage:
- Size: 4.65 MB
- Stars: 1
- Watchers: 1
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: readme.MD
- License: LICENSE
Awesome Lists containing this project
README
# PyTorch Lightning Image Classifiers
This repository contains modular PyTorch Lightning implementations of popular deep learning models for image classification. All models inherit from a common `BaseClassifier` class, making it easy to modify, extend, and use for various tasks.
## Currently Supported Classifiers
| Model | Train | Test | Inference | GRAD-CAM |
|------------------|-------|------|-----------|----------|
| SwinTransformer | ✅ | ✅ | ✅ | ✅ |
| ResNet | ✅ | ✅ | ✅ | ❌ |
| ResNext | ✅ | ✅ | ✅ | ❌ |
| DenseNet | ✅ | ✅ | ✅ | ❌ |
| EfficientNet | ✅ | ✅ | ✅ | ❌ |
| ViT | ✅ | ✅ | ✅ | ❌ |
## Features
- **Modular Design**: All models inherit from `BaseClassifier`, ensuring consistent training, validation, and testing workflows.
- **Easy Configuration**: Modify hyperparameters like learning rate, batch size, and optimizer directly in the configuration.
- **Checkpointing**: Automatically saves the best model during training.
- **Early Stopping**: Prevents overfitting by stopping training if validation performance plateaus.
## Usage
1. **Install Dependencies**:
```bash
pip install pytorch_lightning torchvision transformers efficientnet-pytorch
```
2. **Benchmarking Models**:
Use the `train_all_models.py` script to train and test all models and check what works the best.:
After completion, you will get report in csv as below.
Based on the metric, decide which model is appropriate for your task.
| Model | Test Accuracy | Test Precision | Test Recall | Test F1 | Training Time | Timestamp |
|------------------|-----------------|-----------------|-----------------|-----------------|----------------|----------------|
| ResNext101 | 0.78899 | 0.79964 | 0.78899 | 0.78879 | 1:47:51.260405 | 20250221_094927 |
| ResNet101 | 0.78899 | 0.79463 | 0.78899 | 0.78770 | 1:03:14.367643 | 20250221_113730 |
| SwinTransformer | 0.88073 | 0.88272 | 0.88073 | 0.88030 | 0:22:18.137873 | 20250221_124052 |
| ViT | 0.84404 | 0.84723 | 0.84404 | 0.84415 | 0:22:27.548590 | 20250221_130314 |
| DenseNet121 | 0.85321 | 0.86235 | 0.85321 | 0.85291 | 0:58:01.608439 | 20250221_132545 |
| EfficientNetB7 | 0.84404 | 0.85011 | 0.84404 | 0.84465 | 3:06:57.563459 | 20250221_142354 |
3. **Train, Test and Inference demo**
Use the `demo.ipynb` to follow the whole workflow with a model architecture.
4. **GRAD-CAM Visualization**
The [grad_cam.ipynb](grad_cam.ipynb) notebook provides a detailed workflow for generating Grad-CAM (Gradient-weighted Class Activation Mapping) heatmaps to visualize which regions of an image are most influential for the model's predictions.
## Extending the Repository:
**To add a new model**:
- Create a new Python file under `models` directory. (e.g., new_model_classifier.py).
- Inherit from BaseClassifier and implement the model-specific logic.
- If needed, Add the new model to the `models_to_test` dictionary in [train_all_models.py](train_all_models.py).
## License
This project is licensed under the MIT License. See [LICENSE](LICENSE) for details.