PyTorch CNN Trainer

Train CNN for your task

A simple engine to fine tune CNNs from torchvision and Pytorch Image models from Ross Wightman.

Why This Package ?

It is very annoying to write training loop and training code for CNN training. Also to support all the training features it takes massive time.

Usually we don't need distributed training and it is very uncomfortable to use argparse and get the job done.

This simplifies the training. It provide you a powerful engine.py which can do lot of training functionalities. Also a dataset.py to load dataset in common scenarios.

Note: - Pytorch Trainer is not a distributed training script.

It will work good for single GPU machine for Google Colab / Kaggle.

But for distributed Training you can use the PyTorch Lightning Trainer (soon).

It will train on multiple GPUs just the way lightning supports (soon).

To install

Install torch and torchvision from PyTorch Run the following in terminal.

pip install -q git+git://github.com/oke-aditya/pytorch_cnn_trainer.git

Or from PyPI now !!

pip install pytorch_cnn_trainer

Docs: -

I have provided some example of how to use this trainer in multiple scenarios. Please check examples folder. Some examples that are provided.

• Fine Tune Torchvision models using fit().
• Fine Tune Torchvision models using train_step() and validation_step().
• Fine Tune Ross Wightman's models.
• Quantization Aware Training head only.
• Quantization Aware Training Fully.
• Mixed Precision Training
• Training with Stochastic Weighted Average (SWA)
• LR Finder implementation

Features: -

• Support PyTorch image models (timm) training and transfer learning.
• Quantization Aware training example.
• Early stopping with patience.
• Support torchvision models trainging and transfer learning.
• Support torchvision quantized models transfer learning.
• Support for Mixed Precision Training.
• L2 Norm Gradient Penalty.
• LR Finder Implementation.
• SWA Stochastic weighted Averaging support for training.
• Add Keras Like fit method.
• Sanity Check method.

Hope this repo helps people to train models using transfer learning.

If you like it do give * and tell people about it.