Skip to content

Train and Inference your custom YOLO-NAS model by Pytorch on Windows

License

Notifications You must be signed in to change notification settings

Andrewhsin/YOLO-NAS-pytorch

Repository files navigation



Build, train, and fine-tune production-ready deep learning SOTA vision models Tweet

Version 3 is out! Notebooks have been updated!



Build with SuperGradients


Support various computer vision tasks

Ready to deploy pre-trained SOTA models

YOLO-NAS architecture is out! The new YOLO-NAS delivers state-of-the-art performance with the unparalleled accuracy-speed performance, outperforming other models such as YOLOv5, YOLOv6, YOLOv7 and YOLOv8. Check it out here: YOLO-NAS.

All Computer Vision Models - Pretrained Checkpoints can be found in the Model Zoo

Classification

Semantic Segmentation

Object Detection

Easy to train SOTA Models

Easily load and fine-tune production-ready, pre-trained SOTA models that incorporate best practices and validated hyper-parameters for achieving best-in-class accuracy. For more information on how to do it go to Getting Started

Quick Installation


pip install super-gradients

Prepare Your Dataset

Your custom dataset should be in COCO or YOLO data format.
Data Format:

├── Dataset
│   ├── train
│   │   ├── images
│   │   │   ├── 1.jpg
│   │   │   ├── 2.jpg
|   |   |   ├── ....
│   │   ├── labels
│   │   │   ├── 1.txt
│   │   │   ├── 2.txt
|   |   |   ├── ....
│   ├── val
│   │   ├── images
│   │   │   ├── 1.jpg
│   │   │   ├── 2.png
|   |   |   ├── ....
│   │   ├── labels
│   │   │   ├── 1.txt
│   │   │   ├── 2.txt
|   |   |   ├── ....
│   ├── test
│   │   ├── images
│   │   │   ├── 1.jpg
│   │   │   ├── 50.jpg
|   |   |   ├── ....
│   │   ├── labels
│   │   │   ├── 1.txt
│   │   │   ├── 50.txt
|   |   |   ├── ....

To training custom model using your custom data. You need to create data.yaml Example:

names:
- class1
- class2
- class3
Dir: '/home/user/my_data'  # path to your custom data
images:
  test: test/images
  train: train/images
  val: valid/images
labels:
  test: test/labels
  train: train/labels
  val: valid/labels

Train Your Custom Dataset

You can train your YOLO-NAS model with Single Command Line

Args

-i, --data: path to data.yaml
-n, --name: Checkpoint dir name
-b, --batch: Training batch size
-e, --epoch: number of training epochs.
-j, --worker: Training number of workers
-m, --model: Model type (eg: yolo_nas_s, yolo_nas_m, yolo_nas_l)
-w, --weight: path to pre-trained model weight (default: coco weight)

Example:

python3 train.py --data /dir/dataset/data.yaml --batch 6 --epoch 100 --model yolo_nas_m

Inference

You can Inference your YOLO-NAS model with Single Command Line

Support

  • Video
  • Camera
  • RTSP
Args

-i, --data: path to data.yaml
-m, --model: Model type (eg: yolo_nas_s, yolo_nas_m, yolo_nas_l)
-w, --weight: path to trained model weight
-s, --source: video path/cam-id/RTSP
-c, --conf: model prediction confidence (0<conf<1)
--save: to save video
--hide: hide video window

Example:

python3 inference.py --data /dir/dataset/data.yaml --model yolo_nas_m --weight /runs/train4/ckpt_best.pth --source /test/video.mp4 --conf 0.66           # video
                                                                                                          --source /test/video.mp4 --conf 0.75 --hide    # to save and hide video window
                                                                                                          --source 0 --conf 0.45                         # Camera
                                                                                                          --source 'rtsp://link' --conf 0.25 --save      # save RTSP video stream

What's New - Version 3.1.1 (May 3rd)


Check out SG full release notes.

Coming soon


  • Pre-trained pose estimation model
  • Test Time Augmentations (TTA)
  • Recipe to train DEKR model(convertable to TRT)
  • Key-points Rescoring for Pose estimation
  • LR finder
  • Data analysis tools

Table of Content


Classification

Transfer Learning

Classification Transfer Learning GitHub source


Semantic Segmentation

Quick Start

Segmentation Quick Start GitHub source


Transfer Learning

Segmentation Transfer Learning GitHub source


How to Connect Custom Dataset

Segmentation How to Connect Custom Dataset GitHub source


Object Detection

Transfer Learning

Detection Transfer Learning GitHub source


How to Connect Custom Dataset

Detection How to Connect Custom Dataset GitHub source


How to Predict Using Pre-trained Model

Segmentation, Detection and Classification Prediction

How to Predict Using Pre-trained Model GitHub source


Advanced Features


Post Training Quantization and Quantization Aware Training

Quantization involves representing weights and biases in lower precision, resulting in reduced memory and computational requirements, making it useful for deploying models on devices with limited resources. The process can be done during training, called Quantization aware training, or after training, called post-training quantization. A full tutorial can be found here.

Post Training Quantization and Quantization Aware Training GitHub source

Knowledge Distillation Training

Knowledge Distillation is a training technique that uses a large model, teacher model, to improve the performance of a smaller model, the student model. Learn more about SuperGradients knowledge distillation training with our pre-trained BEiT base teacher model and Resnet18 student model on CIFAR10 example notebook on Google Colab for an easy to use tutorial using free GPU hardware

Knowledge Distillation Training GitHub source


Recipes

To train a model, it is necessary to configure 4 main components. These components are aggregated into a single "main" recipe .yaml file that inherits the aforementioned dataset, architecture, raining and checkpoint params. It is also possible (and recomended for flexibility) to override default settings with custom ones. All recipes can be found here
Recipes support out of the box every model, metric or loss that is implemented in SuperGradients, but you can easily extend this to any custom object that you need by "registering it". Check out this tutorial for more information.

How to Use Recipes GitHub source



Using Distributed Data Parallel (DDP)

Why use DDP ?

Recent Deep Learning models are growing larger and larger to an extent that training on a single GPU can take weeks. In order to train models in a timely fashion, it is necessary to train them with multiple GPUs. Using 100s GPUs can reduce training time of a model from a week to less than an hour.

How does it work ?

Each GPU has its own process, which controls a copy of the model and which loads its own mini-batch from disk and sends it to its GPU during training. After the forward pass is completed on every GPU, the gradient is reduced across all GPUs, yielding to all the GPUs having the same gradient locally. This leads to the model weights to stay synchronized across all GPUs after the backward pass.

How to use it ?

You can use SuperGradients to train your model with DDP in just a few lines.

main.py

from super_gradients import init_trainer, Trainer
from super_gradients.common import MultiGPUMode
from super_gradients.training.utils.distributed_training_utils import setup_device

# Initialize the environment
init_trainer()

# Launch DDP on 4 GPUs'
setup_device(multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL, num_gpus=4)

# Call the trainer
Trainer(expriment_name=...)

# Everything you do below will run on 4 gpus

...

Trainer.train(...)

Finally, you can launch your distributed training with a simple python call.

python main.py

Please note that if you work with torch<1.9.0 (deprecated), you will have to launch your training with either torch.distributed.launch or torchrun, in which case nproc_per_node will overwrite the value set with gpu_mode:

python -m torch.distributed.launch --nproc_per_node=4 main.py
torchrun --nproc_per_node=4 main.py

Calling functions on a single node

It is often in DDP training that we want to execute code on the master rank (i.e rank 0). In SG, users usually execute their own code by triggering "Phase Callbacks" (see "Using phase callbacks" section below). One can make sure the desired code will only be ran on rank 0, using ddp_silent_mode or the multi_process_safe decorator. For example, consider the simple phase callback below, that uploads the first 3 images of every batch during training to the Tensorboard:

from super_gradients.training.utils.callbacks import PhaseCallback, PhaseContext, Phase
from super_gradients.common.environment.env_helpers import multi_process_safe

class Upload3TrainImagesCalbback(PhaseCallback):
    def __init__(
        self,
    ):
        super().__init__(phase=Phase.TRAIN_BATCH_END)
    
    @multi_process_safe
    def __call__(self, context: PhaseContext):
        batch_imgs = context.inputs.cpu().detach().numpy()
        tag = "batch_" + str(context.batch_idx) + "_images"
        context.sg_logger.add_images(tag=tag, images=batch_imgs[: 3], global_step=context.epoch)

The @multi_process_safe decorator ensures that the callback will only be triggered by rank 0. Alternatively, this can also be done by the SG trainer boolean attribute (which the phase context has access to), ddp_silent_mode, which is set to False iff the current process rank is zero (even after the process group has been killed):

from super_gradients.training.utils.callbacks import PhaseCallback, PhaseContext, Phase

class Upload3TrainImagesCalbback(PhaseCallback):
    def __init__(
        self,
    ):
        super().__init__(phase=Phase.TRAIN_BATCH_END)

    def __call__(self, context: PhaseContext):
        if not context.ddp_silent_mode:
            batch_imgs = context.inputs.cpu().detach().numpy()
            tag = "batch_" + str(context.batch_idx) + "_images"
            context.sg_logger.add_images(tag=tag, images=batch_imgs[: 3], global_step=context.epoch)

Note that ddp_silent_mode can be accessed through SgTrainer.ddp_silent_mode. Hence, it can be used in scripts after calling SgTrainer.train() when some part of it should be ran on rank 0 only.

Good to know

Your total batch size will be (number of gpus x batch size), so you might want to increase your learning rate. There is no clear rule, but a rule of thumb seems to be to linearly increase the learning rate with the number of gpus

Easily change architectures parameters

from super_gradients.training import models

# instantiate default pretrained resnet18
default_resnet18 = models.get(model_name="resnet18", num_classes=100, pretrained_weights="imagenet")

# instantiate pretrained resnet18, turning DropPath on with probability 0.5
droppath_resnet18 = models.get(model_name="resnet18", arch_params={"droppath_prob": 0.5}, num_classes=100, pretrained_weights="imagenet")

# instantiate pretrained resnet18, without classifier head. Output will be from the last stage before global pooling
backbone_resnet18 = models.get(model_name="resnet18", arch_params={"backbone_mode": True}, pretrained_weights="imagenet")

Using phase callbacks

from super_gradients import Trainer
from torch.optim.lr_scheduler import ReduceLROnPlateau
from super_gradients.training.utils.callbacks import Phase, LRSchedulerCallback
from super_gradients.training.metrics.classification_metrics import Accuracy

# define PyTorch train and validation loaders and optimizer

# define what to be called in the callback
rop_lr_scheduler = ReduceLROnPlateau(optimizer, mode="max", patience=10, verbose=True)

# define phase callbacks, they will fire as defined in Phase
phase_callbacks = [LRSchedulerCallback(scheduler=rop_lr_scheduler,
                                       phase=Phase.VALIDATION_EPOCH_END,
                                       metric_name="Accuracy")]

# create a trainer object, look the declaration for more parameters
trainer = Trainer("experiment_name")

# define phase_callbacks as part of the training parameters
train_params = {"phase_callbacks": phase_callbacks}

Integration to DagsHub

Open In Colab

from super_gradients import Trainer

trainer = Trainer("experiment_name")
model = ...

training_params = { ...  # Your training params
                   "sg_logger": "dagshub_sg_logger",  # DagsHub Logger, see class super_gradients.common.sg_loggers.dagshub_sg_logger.DagsHubSGLogger for details
                   "sg_logger_params":  # Params that will be passes to __init__ of the logger super_gradients.common.sg_loggers.dagshub_sg_logger.DagsHubSGLogger
                     {
                       "dagshub_repository": "<REPO_OWNER>/<REPO_NAME>", # Optional: Your DagsHub project name, consisting of the owner name, followed by '/', and the repo name. If this is left empty, you'll be prompted in your run to fill it in manually.
                       "log_mlflow_only": False, # Optional: Change to true to bypass logging to DVC, and log all artifacts only to MLflow  
                       "save_checkpoints_remote": True,
                       "save_tensorboard_remote": True,
                       "save_logs_remote": True,
                     }
                   }

Integration to Weights and Biases

from super_gradients import Trainer

# create a trainer object, look the declaration for more parameters
trainer = Trainer("experiment_name")

train_params = { ... # training parameters
                "sg_logger": "wandb_sg_logger", # Weights&Biases Logger, see class WandBSGLogger for details
                "sg_logger_params": # paramenters that will be passes to __init__ of the logger 
                  {
                    "project_name": "project_name", # W&B project name
                    "save_checkpoints_remote": True
                    "save_tensorboard_remote": True
                    "save_logs_remote": True
                  } 
               }

Integration to ClearML

from super_gradients import Trainer

# create a trainer object, look the declaration for more parameters
trainer = Trainer("experiment_name")

train_params = { ... # training parameters
                "sg_logger": "clearml_sg_logger", # ClearML Logger, see class ClearMLSGLogger for details
                "sg_logger_params": # paramenters that will be passes to __init__ of the logger 
                  {
                    "project_name": "project_name", # ClearML project name
                    "save_checkpoints_remote": True,
                    "save_tensorboard_remote": True,
                    "save_logs_remote": True,
                  } 
               }

Installation Methods


Prerequisites

General requirements
To train on nvidia GPUs

Quick Installation

Install stable version using PyPi

See in PyPi

pip install super-gradients

That's it !

Install using GitHub
pip install git+https://github.com/Deci-AI/super-gradients.git@stable

Implemented Model Architectures


All Computer Vision Models - Pretrained Checkpoints can be found in the Model Zoo

Detailed list can be found here

Image Classification

Semantic Segmentation

Object Detection


Implemented Datasets


Deci provides implementation for various datasets. If you need to download any of the dataset, you can find instructions.

Image Classification

Semantic Segmentation

Object Detection


Documentation

Check SuperGradients Docs for full documentation, user guide, and examples.

Contributing

To learn about making a contribution to SuperGradients, please see our Contribution page.

Our awesome contributors:


Made with contrib.rocks.

Citation

If you are using SuperGradients library or benchmarks in your research, please cite SuperGradients deep learning training library.

Community

If you want to be a part of SuperGradients growing community, hear about all the exciting news and updates, need help, request for advanced features, or want to file a bug or issue report, we would love to welcome you aboard!

  • Discord is the place to be and ask questions about SuperGradients and get support. Click here to join our Discord Community

  • To report a bug, file an issue on GitHub.

  • Join the SG Newsletter for staying up to date with new features and models, important announcements, and upcoming events.

  • For a short meeting with us, use this link and choose your preferred time.

License

This project is released under the Apache 2.0 license.

Citing

BibTeX

@misc{supergradients,
  doi = {10.5281/ZENODO.7789328},
  url = {https://zenodo.org/record/7789328},
  author = {Aharon,  Shay and {Louis-Dupont} and {Ofri Masad} and Yurkova,  Kate and {Lotem Fridman} and {Lkdci} and Khvedchenya,  Eugene and Rubin,  Ran and Bagrov,  Natan and Tymchenko,  Borys and Keren,  Tomer and Zhilko,  Alexander and {Eran-Deci}},
  title = {Super-Gradients},
  publisher = {GitHub},
  journal = {GitHub repository},
  year = {2021},
}

Latest DOI

DOI


Deci Platform

Deci Platform is our end to end platform for building, optimizing and deploying deep learning models to production.

Request free trial to enjoy immediate improvement in throughput, latency, memory footprint and model size.

Features

  • Automatically compile and quantize your models with just a few clicks (TensorRT, OpenVINO).
  • Gain up to 10X improvement in throughput, latency, memory and model size.
  • Easily benchmark your models’ performance on different hardware and batch sizes.
  • Invite co-workers to collaborate on models and communicate your progress.
  • Deci supports all common frameworks and Hardware, from Intel CPUs to Nvidia's GPUs and Jetsons. ֿ

Request free trial here

About

Train and Inference your custom YOLO-NAS model by Pytorch on Windows

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published