Pipeline for training Multiclass classifier in PyTorch with Timm model library and TensorRT/Onnx Export

❓ How to use

The script expects images/labels info stored in CSV file and dataset in folder. Check Scene Classification dataset example. Link in reference

Command : python3 train.py --model_name resnet50 --epochs 100 --batch_size 64 --lr 0.0001 --img_size 256 --device 0 --optimizer adam --lr_scheduler CosineAnnealingLR --dataset /home/SceneData --split 0.2 --target_size 3 --early_stop 10 --loss_func CrossEntropyLoss --save_checkpoint_folder ./checkpoints --save_model_folder ./weights --exp_name testExp --labels night,day,noon --wandb --projec_name SceneClassifier --seed 22 --workers 4

If you want to train on dataset which is stored in folders then pass : python3 train.py OTHER_ARGUEMENTS_AS_MENTIONED_ABOVE --image_folder flag along with --train_image_folder trainFolder --val_image_folder validFolder

In this case, dataset structure should be like ->
DatasetName
 |-> train_data
  |-> class_1
  |-> class_2
  |-> class_3
 |-> val_data
  |-> class_1
  |-> class_2
  |-> class_3

🔥 Features/Options/Support

1. Custom model file with Timm models
2. Custom dataset file
3. Creates Experiment folder allowing you to run continous training jobs.
4. Support training for dataset stored in folder or in CSV
5. Categorical Labels - [1], [2], [3]...
6. Displays Training condiguration so you cross check the input
7. Run validation on Train and test set and saves ConfusionMatrix as PNG.
8. Use of garbage collector and torch's method to clear GPU cache
9. Early stopping to save your time and resources
10. Saves model checkpoint and weight file as accuracy improves
11. Albumentations for image augmentation
12. Added options for Loss functions
13. Multiple LR Schedulers
14. Added WandB support
15. Added Number of workers parameter
16. Added Seed to help reproduce experiment
17. Added Torch inference code
18. Added Torch to onnx model export
19. Added Onnx inference code
20. Added Onnx to TensorRT conversion code
21. Added TensorRT inference code

⛳ To/Do

1. Heatmap of features
2. Add LabelSmoothing
3. Add Gradient Clipping
4. Add Mixed Precision Training

💢 Set up your environment

1. Install venv (recommended)
2. Install all the requirements using requirements.txt
3. To install TensorRT, refer Nvidia's Link

♦️ References

1. Scene Classification Dataset
2. Kaggle Notebook : Transfer Learning with Timm
3. Kaggle Notebook : EfficientNet Mixup Leak free
4. Kaggle Notebook : Scene classification
5. Convert PyTorch model to TensorRT - Link
6. Getting started with PyTorch model & Timm - Link

❤️ Made by Sahil Chachra