Skip to content

Latest commit

 

History

History
65 lines (55 loc) · 3.35 KB

README.md

File metadata and controls

65 lines (55 loc) · 3.35 KB
Raw

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