PyTorch CNN Trainer
A simple package to fine-tune CNNs from torchvision and Pytorch Image models by Ross Wightman.
torchvision convolutional-neural-networks pytorch training fine-tuning code computer-vision demo library


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

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 which can do lot of training functionalities. Also a to load dataset in common scenarios.

It has documentation and examples for common to advanced training scenarios too.

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.
  • SWA Stochastic weighted Averaging support for training.
  • Keras Like fit method.
  • Customizable Train Step and Validation Step Method
  • Sanity Check method.

Don't forget to tag @oke-aditya in your comment, otherwise they may not be notified.

Authors original post
Deep Learning | Computer Vision
Share this project
Similar projects
Torchvision Object Detection Finetuning Tutorial
Finetuning a pre-trained Mask R-CNN model in the Penn-Fudan Database for Pedestrian Detection and Segmentation.
Training an Image Classifier in PyTorch
Torchvision, that has data loaders for common datasets such as Imagenet, CIFAR10, MNIST, etc. and data transformers for images, vizualization and data ...
Collection of models and extensions for mobile deployment in PyTorch, built on top of TorchVision and HuggingFace's Transformers.
Top collections