README (README.md) Copy Section

<div align="center" markdown="1">
<img src="documentation/assets/SG_img/SG - Horizontal Glow 2.png" width="600"/>
<br/><br/>

Build, train, and fine-tune production-ready deep learning SOTA vision models
[](https://twitter.com/intent/tweet?text=Easily%20train%20or%20fine-tune%20SOTA%20computer%20vision%20models%20from%20one%20training%20repository&url=https://github.com/Deci-AI/super-gradients&via=deci_ai&hashtags=AI,deeplearning,computervision,training,opensource)

#### Version 3.5 is out! Notebooks have been updated!
______________________________________________________________________
</div>
<div align="center">
<p align="center">
<a href="#getting-started">Getting Started</a> •
<a href="#implemented-model-architectures">Pretrained Models</a> •
<a href="#community">Community</a> •
<a href="#license">License</a> •
</p>
<p align="center">
<a href="https://github.com/Deci-AI/super-gradients#prerequisites"><img src="https://img.shields.io/badge/python-3.7%20%7C%203.8%20%7C%203.9-blue" /></a>
<a href="https://github.com/Deci-AI/super-gradients#prerequisites"><img src="https://img.shields.io/badge/pytorch-1.9%20%7C%201.10-blue" /></a>
<a href="https://pypi.org/project/super-gradients/"><img src="https://img.shields.io/pypi/v/super-gradients" /></a>
<a href="https://github.com/Deci-AI/super-gradients/blob/master/documentation/source/model_zoo.md" ><img src="https://img.shields.io/badge/pre--trained%20models-34-brightgreen" /></a>
<a href="https://github.com/Deci-AI/super-gradients/releases"><img src="https://img.shields.io/github/v/release/Deci-AI/super-gradients" /></a>
<a href="https://join.slack.com/t/supergradients-comm52/shared_invite/zt-10vz6o1ia-b_0W5jEPEnuHXm087K~t8Q"><img src="https://img.shields.io/badge/slack-community-blueviolet" /></a>
<a href="https://github.com/Deci-AI/super-gradients/blob/master/LICENSE.md"><img src="https://img.shields.io/badge/license-Apache%202.0-blue" /></a>
<a href="https://docs.deci.ai/super-gradients/documentation/source/welcome.html"><img src="https://img.shields.io/badge/docs-mkdocs-brightgreen" /></a>
</p>
</div>

______________________________________________________________________

Build with SuperGradients Copy Chapter

Support various computer vision tasks Copy Section

Ready to deploy pre-trained SOTA models Copy Section

YOLO-NAS and YOLO-NAS-POSE architectures are 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.
A YOLO-NAS-POSE model for pose estimation is also available, delivering state-of-the-art accuracy/performance tradeoff.

Check these out here: YOLO-NAS & YOLO-NAS-POSE.

<div align="center">
<img src="https://github.com/Deci-AI/super-gradients/raw/master/documentation/source/images/yolo_nas_frontier.png" height="600px">
<img src="https://github.com/Deci-AI/super-gradients/raw/master/documentation/source/images/yolo_nas_pose_frontier_t4.png" height="600px">
</div>

Load model with pretrained weights

from super_gradients.training import models
from super_gradients.common.object_names import Models
model = models.get(Models.YOLO_NAS_M, pretrained_weights="coco")

#### Classification
<div align="center">
<img src="https://github.com/Deci-AI/super-gradients/raw/master/documentation/assets/SG_img/[email protected]" width="800px">
</div>

#### Semantic Segmentation
<div align="center">
<img src="https://github.com/Deci-AI/super-gradients/raw/master/documentation/assets/SG_img/Semantic [email protected]" width="800px">
</div>

#### Object Detection
<div align="center">
<img src="https://github.com/Deci-AI/super-gradients/raw/master/documentation/assets/SG_img/Object [email protected]" width="800px">
</div>

#### Pose Estimation
<div align="center">
<img src="https://github.com/Deci-AI/super-gradients/raw/master/documentation/source/images/yolo_nas_pose_frontier_t4.png" width="400px">
<img src="https://github.com/Deci-AI/super-gradients/raw/master/documentation/source/images/yolo_nas_pose_frontier_xavier_nx.png" width="400px">
</div>

Easy to train SOTA Models Copy Section

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

#### Plug and play recipes

python -m super_gradients.train_from_recipe architecture=regnetY800 dataset_interface.data_dir=<YOUR_Imagenet_LOCAL_PATH> ckpt_root_dir=<CHEKPOINT_DIRECTORY>

Production readiness Copy Section

Load model with pretrained weights

from super_gradients.training import models
from super_gradients.common.object_names import Models
model = models.get(Models.YOLO_NAS_M, pretrained_weights="coco")
Prepare model for conversion

Input size is in format of [Batch x Channels x Width x Height] where 640 is the standard COCO dataset dimensions

model.eval()
model.prep_model_for_conversion(input_size=[1, 3, 640, 640])
    
Create dummy_input
Convert model to onnx

torch.onnx.export(model, dummy_input,  "yolo_nas_m.onnx")

Quick Installation Copy Chapter

__________________________________________________________________________________________________________

pip install super-gradients

What's New Copy Chapter

__________________________________________________________________________________________________________
Version 3.4.0 (November 6, 2023)

* YoloNAS-Pose model released - a new frontier in pose estimation
* Added option to export a recipe to a single YAML file or to a standalone train.py file
* Other bugfixes & minor improvements. Full release notes available here

__________________________________________________________________________________________________________
Version 3.1.3 (July 19, 2023)

* Pose Estimation Task Support - Check out fine-tuning notebook example
* Pre-trained modified DEKR model for pose estimation (TensorRT-compatible)
* Support for Python 3.10
* Support for torch.compile
* Other bugfixes & minor improvements. Check out release notes

__________________________________________________________________________________________________________
30th of May
* Quantization Aware Training YoloNAS on Custom Dataset

Version 3.1.1 (May 3rd)
* YOLO-NAS
* New predict function (predict on any image, video, url, path, stream)
* RoboFlow100 datasets integration
* A new Documentation Hub
* Integration with DagsHub for experiment monitoring
* Support Darknet/Yolo format detection dataset (used by Yolo v5, v6, v7, v8)
* Segformer model and recipe
* Post Training Quantization and Quantization Aware Training - notebooks

Check out SG full release notes.

Table of Content Copy Chapter

- Getting Started
- Advanced Features
- Installation Methods
- Prerequisites
- Quick Installation
- Implemented Model Architectures
- Contributing
- Citation
- Community
- License
- Deci Platform

Getting Started Copy Chapter

Start Training with Just 1 Command Line Copy Section

Just make sure that you setup your dataset according to the data dir specified in the recipe.

python -m super_gradients.train_from_recipe --config-name=imagenet_regnetY architecture=regnetY800 dataset_interface.data_dir=<YOUR_Imagenet_LOCAL_PATH> ckpt_root_dir=<CHEKPOINT_DIRECTORY>

Quickly Load Pre-Trained Weights for Your Desired Model with SOTA Performance Copy Section

The pretrained_weights argument will load a pre-trained architecture on the provided dataset

    
import super_gradients
model = models.get("model-name", pretrained_weights="pretrained-model-name")

Classification Copy Section

* [](https://colab.research.google.com/github/Deci-AI/super-gradients/blob/master/notebooks/transfer_learning_classification.ipynb) Transfer Learning for classification
* [](https://colab.research.google.com/github/Deci-AI/super-gradients/blob/master/notebooks/PTQ_and_QAT_for_classification.ipynb) PTQ and QAT for classification

Semantic Segmentation Copy Section

* [](https://colab.research.google.com/github/Deci-AI/super-gradients/blob/master/notebooks/quickstart_segmentation.ipynb) Segmentation Quick Start
* [](https://colab.research.google.com/github/Deci-AI/super-gradients/blob/master/notebooks/transfer_learning_semantic_segmentation.ipynb) Segmentation Transfer Learning
* [](https://colab.research.google.com/github/Deci-AI/super-gradients/blob/master/notebooks/segmentation_connect_custom_dataset.ipynb) How to Connect Custom Dataset
* [](https://colab.research.google.com/github/Deci-AI/super-gradients/blob/master/notebooks/Segmentation_Model_Export.ipynb) How to export segmentation model to ONNX

Pose Estimation Copy Section

* [](https://colab.research.google.com/github/Deci-AI/super-gradients/blob/master/notebooks/YoloNAS_Pose_Fine_Tuning_Animals_Pose_Dataset.ipynb) Fine Tuning YoloNAS-Pose on AnimalPose dataset
* [](https://colab.research.google.com/github/Deci-AI/super-gradients/blob/master/notebooks/DEKR_PoseEstimationFineTuning.ipynb) Fine Tuning DEKR on AnimalPose dataset

Object Detection Copy Section

* [](https://colab.research.google.com/github/Deci-AI/super-gradients/blob/master/notebooks/YoloNAS_Inference_using_TensorRT.ipynb) YoloNAS inference using TensorRT
* [](https://colab.research.google.com/github/Deci-AI/super-gradients/blob/master/notebooks/detection_transfer_learning.ipynb) Object Detection Transfer Learning
* [](https://colab.research.google.com/github/Deci-AI/super-gradients/blob/master/notebooks/detection_how_to_connect_custom_dataset.ipynb) How to Connect Custom Dataset
* [](https://colab.research.google.com/github/Deci-AI/super-gradients/blob/master/notebooks/yolo_nas_custom_dataset_fine_tuning_with_qat.ipynb) Quantization Aware Training YoloNAS on Custom Dataset

How to Predict Using Pre-trained Model Copy Section

* [](https://colab.research.google.com/github/Deci-AI/super-gradients/blob/master/notebooks/how_to_run_model_predict.ipynb) How to Predict Using Pre-trained Model

Albumentations Integration Copy Section

* [](https://colab.research.google.com/github/Deci-AI/super-gradients/blob/master/notebooks/albumentations_tutorial.ipynb) Using Albumentations with SG

Advanced Features Copy Chapter

Post Training Quantization and Quantization Aware Training Copy Section

* [](https://bit.ly/3KrN6an) Post Training Quantization and Quantization Aware Training

Quantization Aware Training YoloNAS on Custom Dataset Copy Section

* [](https://bit.ly/3MIKdTy) Quantization Aware Training YoloNAS on Custom Dataset

Knowledge Distillation Training Copy Section

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

* [](https://colab.research.google.com/github/Deci-AI/super-gradients/blob/master/notebooks/how_to_use_knowledge_distillation_for_classification.ipynb) Knowledge Distillation Training

Recipes Copy Section

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 recommended for flexibility) to override default settings with custom ones.
All recipes can be found here
</br>
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.

* [](https://colab.research.google.com/github/Deci-AI/super-gradients/blob/master/notebooks/what_are_recipes_and_how_to_use.ipynb) How to Use Recipes

<details markdown="1">
<summary><h3>Using Distributed Data Parallel (DDP) </h3></summary>

#### 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)

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

</details>

<details markdown="1">
<summary><h3> Easily change architectures parameters </h3></summary>

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")

</details>

<details markdown="1">

<summary><h3> Using phase callbacks </h3></summary>

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}

</details>

<details markdown="1">

<summary><h3> Integration to DagsHub </h3></summary>

[](https://colab.research.google.com/drive/11fW56pMpwOMHQSbQW6xxMRYvw1mEC-t-?usp=sharing)

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,
                     }
                   }

</details>

<details>

<summary><h3> Integration to Weights and Biases </h3></summary>

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
                  } 
               }

</details>

<details markdown="1">

<summary><h3> Integration to ClearML </h3></summary>

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,
                  } 
               }

</details>
<details markdown="1">

<summary><h3> Integration to Voxel51 </h3></summary>

You can apply SuperGradients YOLO-NAS models directly to your FiftyOne dataset using the apply_model() method:

import fiftyone as fo
import fiftyone.zoo as foz
from super_gradients.training import models
dataset = foz.load_zoo_dataset("quickstart", max_samples=25)
dataset.select_fields().keep_fields()
model = models.get("yolo_nas_m", pretrained_weights="coco")
dataset.apply_model(model, label_field="yolo_nas", confidence_thresh=0.7)
session = fo.launch_app(dataset)

The SuperGradients YOLO-NAS model can be accessed directly from the FiftyOne Model Zoo:

import fiftyone as fo
import fiftyone.zoo as foz
model = foz.load_zoo_model("yolo-nas-torch")
dataset = foz.load_zoo_dataset("quickstart")
dataset.apply_model(model, label_field="yolo_nas")
session = fo.launch_app(dataset)

</details>

Installation Methods Copy Chapter

Prerequisites Copy Section

</details>

<details markdown="1">

<summary>To train on nvidia GPUs</summary>

- Nvidia CUDA Toolkit >= 11.2
- CuDNN >= 8.1.x
- Nvidia Driver with CUDA >= 11.2 support (≥460.x)

</details>

Quick Installation Copy Section

<details markdown="1">

<summary>Install stable version using PyPi</summary>

See in PyPi

pip install super-gradients

That's it !

</details>

<details markdown="1">

<summary>Install using GitHub</summary>

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

</details>

Implemented Model Architectures Copy Chapter

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

Image Classification Copy Section

Semantic Segmentation Copy Section

- PP-LiteSeg
- DDRNet (Deep Dual-resolution Networks)
- LadderNet
- RegSeg
- ShelfNet
- STDC

Object Detection Copy Section

Pose Estimation Copy Section

- DEKR-W32-NO-DC

__________________________________________________________________________________________________________

Implemented Datasets Copy Chapter

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

Image Classification Copy Section

Semantic Segmentation Copy Section

- Cityscapes
- Coco
- PascalVOC 2012 / PascalAUG 2012
- SuperviselyPersons
- Mapillary Vistas Dataset

Object Detection Copy Section

Pose Estimation Copy Section

- COCO

__________________________________________________________________________________________________________

Documentation Copy Chapter

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

Contributing Copy Chapter

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

Our awesome contributors:

<a href="https://github.com/Deci-AI/super-gradients/graphs/contributors">
<img src="https://contrib.rocks/image?repo=Deci-AI/super-gradients" />
</a>

<br/>Made with contrib.rocks.

Citation Copy Chapter

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

Community Copy Chapter

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 Copy Chapter

This project is released under the Apache 2.0 license.

Citing Copy Chapter

BibTeX Copy Section

@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 Copy Section

[](https://doi.org/10.5281/zenodo.7789328)

__________________________________________________________________________________________________________

ֿ

Request free trial here

---

CONTRIBUTING (CONTRIBUTING.md) Copy Section

Contribution Guidelines

Here is a simple guideline to get you started with your first contribution.
1. Set up your environment to follow our formatting guidelines and to use signed-commits.
2. Use issues to discuss the suggested changes. Create an issue describing changes if necessary and add labels to ease orientation.
3. Fork super-gradients so you can make local changes and test them.
4. Create a new branch for the issue. The branch naming convention is enforced by the CI/CD so please make sure you are using feature/SG- or hotfix/SG- format otherwise it will fail.
5. Implement your changes along with relevant tests for the issue. Please make sure you are covering unit, integration and e2e tests where required.
6. Create a pull request against master branch.

Code Style Copy Chapter

We follow the reStructuredText docstring format (default of PyCharm), along with typing.

def python_function(first_argument: int, second_argument: int) -> str:
    """Do something with the two arguments.
    :param first_argument: First argument to the function
    :param second_argument: Second argument to the function
    :return: Description of the output
    """

Code Formatting Copy Chapter

We enforce black code formatting in addition to existing flake8 checks.

To ensure everyone uses same code style, a project-wise configuration file has been added to SG repo. It ensures all formatting will be exactly the same regardless of OS, python version or the place where code formatting check is happening.

Installation Copy Section

To start, one need to install required development dependencies (actual versions of black, flake8 and git commit hooks):

$ pip install -r requirements.dev.txt

Pre-Commit Hooks Copy Section

A pre-commit hook as an easy way to ensure all files in the commit are already formatted accordingly and pass linter checks. If they are not, the git will prevent commit of problematic files unless errors are fixed.

To start, run the following command from SG repo root:

$ pip install pre-commit
$ pre-commit install

The command should complete without errors. Once done, all your upcoming commits will be checked via black & flake8.

Usage Copy Section

Just run ``$ black .` from the SG root. It will reformat the whole repo.
For flake8: `$ flake8 --statistics --config scripts/flake8-config setup.py .`

Signed Commits Copy Chapter

Background Copy Section

Signed commits provide a way to verify the authenticity and integrity of the code changes made by a particular developer, as the commit is cryptographically signed using their private GPG key. This helps ensure that the code changes were made by the intended person and have not been tampered with during transit.

You can find more information here.

Add GPG key to GitHub Copy Section

1. Generate a new GPG key. After completing the steps 1-9 you'll see the console message containing generated key ID (e.g. <i>3AA5C34371567BD2</i>).
2. Copy the GPG key by running the command on step 12 from the link above



$ gpg --armor --export <YOUR_GPG_KEY_ID>

$ git config --global user.signingkey <YOUR_GPG_KEY_ID>
$ git config --global gpg.program $(which gpg)
$ git config --global commit.gpgsign true

git commit -S -m "Making my first signed commit"

git log --show-signature -1

git checkout master
git checkout -b feature/my_awesome_pr_signed
git merge --squash --no-commit feature/my_awesome_pr
git commit -S -m "My signed commit"
git push -f origin feature/my_awesome_pr_signed:feature/my_awesome_pr

error: gpg failed to sign the data
fatal: failed to write commit object

18:57:12.099725 run-command.c:663       trace: run_command: /usr/bin/gpg --status-fd=2 -bsau <YOUR_GPG_KEY_ID>

pip install nbdime

nbdiff notebook_1.ipynb notebook_2.ipynb

import super_gradients
yolo_nas = super_gradients.training.models.get("yolo_nas_pose_l", pretrained_weights="coco_pose").cuda()
model_predictions  = yolo_nas.predict("https://deci-pretrained-models.s3.amazonaws.com/sample_images/beatles-abbeyroad.jpg", conf=0.5).show()
prediction = model_predictions[0].prediction # One prediction per image - Here we work with 1 image, so we get the first.
bboxes = prediction.bboxes_xyxy # [Num Instances, 4] List of predicted bounding boxes for each object 
poses  = prediction.poses       # [Num Instances, Num Joints, 3] list of predicted joints for each detected object (x,y, confidence)
scores = prediction.scores      # [Num Instances] - Confidence value for each predicted instance

import super_gradients
yolo_nas = super_gradients.training.models.get("yolo_nas_l", pretrained_weights="coco").cuda()
model_predictions  = yolo_nas.predict("https://deci-pretrained-models.s3.amazonaws.com/sample_images/beatles-abbeyroad.jpg").show()
prediction = model_predictions[0].prediction        # One prediction per image - Here we work with 1 image so we get the first.
bboxes = prediction.bboxes_xyxy                     # [[Xmin,Ymin,Xmax,Ymax],..] list of all annotation(s) for detected object(s) 
bboxes = prediction.bboxes_xyxy                     # [[Xmin,Ymin,Xmax,Ymax],..] list of all annotation(s) for detected object(s) 
class_names = prediction.class_names                # ['Class1', 'Class2', ...] List of the class names
class_name_indexes = prediction.labels.astype(int)  # [2, 3, 1, 1, 2, ....] Index of each detected object in class_names(corresponding to each bounding box)
confidences =  prediction.confidence.astype(float)  # [0.3, 0.1, 0.9, ...] Confidence value(s) in float for each bounding boxes

from super_gradients import Trainer
trainer = Trainer("experiment_name")
model = ...
training_params = {"mixed_precision": True, ...:...}
trainer.train(model=model, training_params=training_params, ...)

my_training_hyperparams.yaml
mixed_precision: True # Whether to use mixed precision or not.

The first step is to export our YoloNAS model to ONNX correctly. Two actions must be taken before we export our model to onnx:
1. We must call model.prep_model_for_conversion - this is essential as YoloNAS incorporates QARepVGG blocks. Without this call, the RepVGG branches will not be fused, and our model's speed will decrease significantly! This is true for the Pytorch model as well as the compiled TRT Engine!
Nothing to worry about if you have quantized your model with PTQ/QAT with SG, as this is done under the hood before exporting the ONNX checkpoints.


2. We need to replace our layers with "fake quantized" ones - this happens when we perform post-training quantization or quantization-aware training with SG.
Again, nothing to worry about if you performed PTQ/QAT with SG and hold your newly exported ONNX checkpoint. Beware that inference time in Pytorch is slower with such blocks - but will be faster once converted to the TRT Engine.

There are plenty of guides on how to perform PTQ/QAT with SG:
- Quantization-aware fine-tuning YoloNAS on custom dataset notebook
- QA/PTQ YoloNAS with configuration files
- QA/PTQ

Suppose we ran PTQ/QAT, then our PTQ/QAT checkpoints have been exported to our checkpoints directory.
If we plug them into netron, we can see that new blocks that were not a part of the original network were introduced: the Quantize/Dequantize layers -

<div>
<img src="images/qdq_yolonas_netron.png" width="750">
</div>

This is expected and an excellent way to verify that our model is ready to be converted to Int8 using Nvidia's TesnorRT.
As stated earlier - inference time in Pytorch is slower with such blocks - but will be faster once converted to the TRT Engine.

Step 2: Create TRT Engine Copy Chapter

trtexec --fp16 --int8 --avgRuns=100 --onnx=your_yolonas_qat_model.onnx

trtexec --fp16 --int8 --avgRuns=100 --onnx=your_yolonas_qat_model.onnx

Command notes:
- Note that this process might take some time, depending on the GPU, the model, and the size of the input (up to 40 minutes is reasonable on smaller devices).
- Also notice the --avgRuns=100 which means that this command runs the model 100 times, so that we get more "robust" results that are less affected by noise.
- Note that since we ran PTQ/QAT we need the --int8 flag flag. But if we did not do so, then the --int8 flag will dramatically degrade the accuracy of our compiled model.
- By passing --fp16 and --int together, we allow the hybrid quantization (that is, some layers are quantized while others are not).

Benchmark breakdown:
- The actual throughput and latency of your model are in blue. This tells you how your model is actually performing.
- Note that trtexec shows the minimum, maximum, mean, and median values. Significant differences between these values can indicate that your measurements are noisy. It might be that some other process uses the GPU while benchmarking or that the GPU needs to be adequately cooled.
- The end-to-end latency is marked in yellow. This includes the time it takes to prepare the input and pass it to the GPU, the GPU compute time, and the time it takes to move the output from the GPU back to the host (for the full batch size). If you plan on running batches one by one synchronously, this is the time that affects you. But if you use an async inference engine, you will be affected by the numbers in blue.
- High H2D (H=Host=CPU; D=Device=GPU) values indicate your input size has a crucial effect on the performance. Consider resizing the input in advance (not on the GPU), or test with different batch sizes to find the optimal setting.
- High D2H values indicate your output might be too big. You can consider a task-specific method to reduce it. i.e., use top-k at the end of your detection model to limit the number of boxes coming out. Alternatively, use a softmax layer at the end of your segmentation model to change the output representation to one with smaller dimensions.

---

Documentation/Source/Checkpoints (documentation/source/Checkpoints.md) Copy Section

Model Checkpoints

If you are not familiar on how experiments are managed, you can check this tutorial


The first question that arises is: what is a checkpoint?

From the Pytorch Lightning documentation:


When a model is training, the performance changes as it sees more data. It is a best practice to save the state of a model throughout the training process. This gives you a version of the model, a checkpoint, at each key point during the development of the model. Once training has been completed, use the checkpoint corresponding to the best performance you found during the training process.

Checkpoints also enable your training to resume where it was in case the training process is interrupted.

Checkpoints Saving: Which, When, and Where? Copy Chapter

From the previous subsection, we understand that checkpoints saved at different times during training have different purposes.
That's why in SG, multiple checkpoints are saved throughout training:

| Checkpoint Filename | When is it saved?|
| ------------- |:-------------:|
| ckpt_best.pth | Each time we reach a new best metric_to_watch when performing validation. |
| ckpt_latest.pth | At the end of every epoch, constantly overriding. |
| average_model.pth | At the end of training - composed of 10 best models according to metric_to_watch and will only be save when the training_param average_best_models=True. |
| ckpt_epoch_{EPOCH_INDEX}.pth | At the end of a fixed epoch number EPOCH_INDEX if it is specified through save_ckpt_epoch_list training_param |

#### Where are the checkpoint files saved?

The checkpoint files will be saved at <ckpt_root_dir>/<experiment_name>/<run_dir>.

- ckpt_root_dir and experiment_name can be set by the user when instantiating the Trainer.


Trainer(ckpt_root_dir='path/to/ckpt_root_dir', experiment_name="my_experiment")

When working with a cloned version of SG, one can leave out the ckpt_root_dir arg, and checkpoints will be saved under the super_gradients/checkpoints directory.

Checkpoint Structure Copy Chapter

Checkpoints in SG are instances of state_dict.
They hold additional information about the model's training besides the model weights.

The checkpoint keys:

- net: The network's state_dict (state_dict).
- acc: The network's achieved metric value on the validation set (metric_to_watch in training_params (float).
- epoch: The last epoch performed.
- optimizer_state_dict: The state_dict of the optimizer (state_dict).
- scaler_state_dict: Optional - only present when training with mixed_precision=True. The state_dict of Trainer.scaler.
- ema_net: Optional - only present when training with ema=True. The EMA model's state_dict. Note that average_model.pth lacks this entry even if ema=True since the average model's snapshots are of the EMA network already (i.e., the "net" entry is already an average of the EMA snapshots).
- torch_scheduler_state_dict: Optional, will only be present when using a torch native lr scheduler (see LRScheduling)


Remote Checkpoint Saving with SG Loggers Copy Chapter

SG supports remote checkpoint saving using 3rd party tools (for example, Weights & Biases).
To do so, specify save_checkpoints_remote=True inside sg_logger_params training_param.
See our documentation on Third-party experiment monitoring.

Loading Checkpoints Copy Chapter

When discussing checkpoint loading in SG, we must separate it into two use cases: loading weights and resuming training.
While the former requires the model's state_dict alone, SG checkpoint loading methods introduce more functionality than PyTorch's vanilla load_state_dict(), especially for SG-trained checkpoints.

Loading Model Weights from a Checkpoint Copy Section

Loading model weights can be done right after model initialization, using models.get(...), or by explicitly calling load_checkpoint_to_model on the torch.nn.Module instance.

Suppose we have launched a training experiment with a similar structure to the one below:

from super_gradients.training import Trainer
...
...
from super_gradients.training import models
from super_gradients.common.object_names import Models
trainer = Trainer("my_resnet18_training_experiment", ckpt_root_dir="/path/to/my_checkpoints_folder")
train_dataloader = ...
valid_dataloader = ...
model = models.get(model_name=Models.RESNET18, num_classes=10)
train_params = {
    ...
    "loss": "CrossEntropyLoss",
    "criterion_params": {},
    "save_ckpt_epoch_list": [10,15]
    ...
}
trainer.train(model=model, training_params=train_params, train_loader=train_dataloader, valid_loader=valid_dataloader)

Then at the end of the training, our ckpt_root_dir contents will look similar to the following:

<ckpt_root_dir>
│
├── <experiment_name>
│   │
│   ├─── <run_dir>
│   │     ├─ ckpt_best.pth                   # Best performance during validation
│   │     ├─ ckpt_latest.pth                 # End of the most recent epoch
│   │     ├─ average_model.pth               # Averaged over specified epochs
│   │     ├─ ckpt_epoch_*.pth                # Checkpoints from specific epochs (like epoch 10, 15, etc.)
│   │     ├─ events.out.tfevents.*           # Tensorflow run artifacts
│   │     └─ log_<timestamp>.txt             # Trainer logs of the specific run
│   │
│   └─── <other_run_dir>
│        └─ ...
│
└─── <other_experiment_name>
    │
    ├─── <run_dir>
    │     └─ ...
    │
    └─── <another_run_dir>
          └─ ...

Suppose we wish to load the weights from ckpt_best.pth. We can simply pass its path to the checkpoint_path argument in models.get(...):

from super_gradients.training import models
from super_gradients.common.object_names import Models
model = models.get(model_name=Models.RESNET18, num_classes=10, checkpoint_path="/path/to/my_checkpoints_folder/my_resnet18_training_experiment/RUN_20230802_131052_651906/ckpt_best.pth")

Important: when loading SG-trained checkpoints using models.get(...), if the network was trained with EMA, the EMA weights will be the ones loaded.

If we already possess an instance of our model, we can also directly use load_checkpoint_to_model:

from super_gradients.training import models
from super_gradients.common.object_names import Models
from super_gradients.training.utils.checkpoint_utils import load_checkpoint_to_model
model = models.get(model_name=Models.RESNET18, num_classes=10)
load_checkpoint_to_model(net=model, ckpt_local_path="/path/to/my_checkpoints_folder/my_resnet18_training_experiment/RUN_20230802_131052_651906/ckpt_best.pth")

Extending the Functionality of PyTorch's strict Parameter in load_state_dict() Parameter in load_state_dict()')" title="Copy section prompt for LLMs"> Copy Section

When not familiar with PyTorch's strict parameter in load_state_dict(), please see PyTorch's docs on this matter first.

The equivalent arguments for PyTorch's strict parameter in load_state_dict() in models.get() and load_checkpoint_to_model are strict and strict_load respectively, and expect SG's StrictLoad enum type.

Let's have a look at its possible values:

class StrictLoad(Enum):
    """
    Wrapper for adding more functionality to torch's strict_load parameter in load_state_dict().
    Attributes:
        OFF              - Native torch "strict_load = off" behavior. See nn.Module.load_state_dict() documentation for more details.
        ON               - Native torch "strict_load = on" behavior. See nn.Module.load_state_dict() documentation for more details.
        NO_KEY_MATCHING  - Allows the usage of SuperGradient's adapt_checkpoint function, which loads a checkpoint by matching each
                           layer's shapes (and bypasses the strict matching of the names of each layer (i.e., disregards the state_dict key matching)).
    """
    OFF = False
    ON = True
    NO_KEY_MATCHING = "no_key_matching"

In other words, we added another loading mode option- no_key_matching. This option exploits the fact that the state_dicts are OrderedDicts, and comes in handy when the underlying network's structure remains the same, but the state_dicts keys do not match the ones inside the models state_dict.
Let's demonstrate the different strict modes with a simple example:

import torch
class ModelA(torch.nn.Module):
    def __init__(self):
        super(ModelA, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 6, 5)
        self.conv2 = torch.nn.Conv2d(6, 16, 5)
class ModelB(torch.nn.Module):
    def __init__(self):
        super(ModelB, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 6, 5)
        self.CONV2 = torch.nn.Sequential([torch.nn.Conv2d(6, 16, 5)])

Notice the above networks have identical weight structures but will have different keys in their state_dicts.
This is why loading a checkpoint from either one to the other, using strict=True, will fail and crash. Using strict=False will not crash and successfully load the first layer's weights only.
Using SG's no_key_matching will successfully load a checkpoint from either one to the other.

Loading Pretrained Weights from the Model Zoo Copy Section

Using models.get(...), you can load any of our pre-trained models in 3 lines of code:


from super_gradients.training import models
from super_gradients.common.object_names import Models
model = models.get(Models.YOLOX_S, pretrained_weights="coco")

The pretrained_weights argument specifies the dataset on which the pre-trained weights were trained. Here is the complete list of pre-trained weights.

Loading Checkpoints: Training with Configuration Files Copy Section

Prerequisites: Training with Configuration Files

Recall the SGs recipes library structure:


The super_gradients/recipes include the following subdirectories

- arch_params - containing configuration files for instantiating different models
- checkpoint_params - containing configuration files that define the loaded and saved checkpoints parameters for the training
- conversion_params - containing configuration files for the model conversion scripts (for deployment)
- dataset_params - containing configuration files for instantiating different datasets and dataloaders
- training_hyperparams - containing configuration files holding hyper-parameters for specific recipes

And now, let's take a look at the default parameters in checkpoint_params:

load_backbone: False # whether to load only the backbone part of the checkpoint
checkpoint_path: # checkpoint path that is located in super_gradients/checkpoints
strict_load: True # key matching strictness for loading checkpoint's weights
pretrained_weights: # a string describing the dataset of the pre-trained weights (for example, "imagenent").

And note the above parameters are used to start the training with different weights (fine-tuning etc.) - they are passed to model.get() in the underlying flow of Trainer.train_from_config(...):


@classmethod
def train_from_config(cls, cfg: Union[DictConfig, dict]) -> Tuple[nn.Module, Tuple]:
    ...
    # BUILD NETWORK
    model = models.get(
        ...
        strict_load=cfg.checkpoint_params.strict_load,
        pretrained_weights=cfg.checkpoint_params.pretrained_weights,
        checkpoint_path=cfg.checkpoint_params.checkpoint_path,
        load_backbone=cfg.checkpoint_params.load_backbone,
    )
    # INSTANTIATE DATA LOADERS
    train_dataloader = ...
    val_dataloader = ...
    
    ...
    # TRAIN
    res = trainer.train(...)
...

Resuming Training Copy Chapter

Resuming training in SG is a comprehensive process, controlled by three primary parameters that allow flexibility
in continuing or branching off from specific training checkpoints.
These parameters are used within training_params: resume, run_id, and resume_path.

resume: False  # Option to continue training from the latest checkpoint.
run_id:        # ID to resume from a specific run within the same experiment.
resume_path:   # Direct path to a specific checkpoint file (.pth) to resume training.
...


#### 1. Resuming the Latest Run

By setting resume=True, SuperGradients will resume training from the last checkpoint within the same experiment.

Example:

Continues from the latest run in the cifar_experiment.

python -m super_gradients.train_from_recipe --config-name=cifar10_resnet experiment_name=cifar_experiment training_hyperparams.resume=True

#### 2. Resuming a Specific Run

Using run_id, you can resume training from a specific run within the same experiment, identified by the run ID.

Example:

Continues from a specific run identified by the ID within cifar_experiment.

python -m super_gradients.train_from_recipe --config-name=cifar10_resnet experiment_name=cifar_experiment run_id=RUN_20230802_131052_651906

#### 3. Branching off from a specific checkpoint

By specifying a resume_path, SuperGradients will create a new run directory,
allowing training to resume from that specific checkpoint,
and subsequently save the new checkpoints in this new directory.

Example:

Branches from a specific checkpoint, creating a new run.

python -m super_gradients.train_from_recipe --config-name=cifar10_resnet experiment_name=cifar_experiment training_hyperparams.resume_path=/path/to/checkpoint.pth

#### 4. Resuming with original recipe
Resuming is parameter dependant - you cannot resume the training of a model if there is a mismatch between the model architecture defined in your recipe, and the one in your checkpoint.

Therefore, if you trained a model a while ago, and that in the meantime you changed the model architecture definition,
then you won't be able to resume its training, loading the model would simply raise an exception.


To avoid this issue, SuperGradients provides an option to resume a training based on the recipe that was originally used to train the model.

Trainer.resume_experiment(ckpt_root_dir=..., experiment_name=..., run_id=...)

- run_id is optional. You can use it to chose which run you want to resume. By default, it will resume the latest run of your experiment.

Note that Trainer.resume_experiment can only resume training that were launched with Trainer.train_from_config.

See usage in our resume_experiment_example.

Resuming Training from SG Logger's Remote Storage (WandB only) Copy Chapter

SG supports saving checkpoints throughout the training process in the remote storage defined by SG Logger (more info about this object and it's role during training in SG at Third-party experiment monitoring.)
Suppose we run an experiment with a WandB SG logger, then our training_hyperparams should hold:


sg_logger: wandb_sg_logger, # Weights&Biases Logger, see class super_gradients.common.sg_loggers.wandb_sg_logger.WandBSGLogger for details
sg_logger_params:             # Params that will be passes to __init__ of the logger super_gradients.common.sg_loggers.wandb_sg_logger.WandBSGLogger
  project_name: project_name, # W&B project name
  save_checkpoints_remote: True,
  save_tensorboard_remote: True,
  save_logs_remote: True,
  entity: <YOUR-ENTITY-NAME>,         # username or team name where you're sending runs
  api_server: <OPTIONAL-WANDB-URL>    # Optional: In case your experiment tracking is not hosted at wandb servers

The save_checkpoints_remote flag is set which will result in saving checkpoints in WandB throughout training.
Now, in case the training was interrupted, we can resume it from the checkpoint located in the WandB run storage by setting 2 training hyperparameters:
1. Set resume_from_remote_sg_logger:


resume_from_remote_sg_logger: True

sg_logger: wandb_sg_logger, # Weights&Biases Logger, see class super_gradients.common.sg_loggers.wandb_sg_logger.WandBSGLogger for details
sg_logger_params:             # Params that will be passes to __init__ of the logger super_gradients.common.sg_loggers.wandb_sg_logger.WandBSGLogger
  wandb_id: <YOUR_RUN_ID>
  project_name: project_name, # W&B project name
  save_checkpoints_remote: True,
  save_tensorboard_remote: True,
  save_logs_remote: True,
  entity: <YOUR-ENTITY-NAME>,         # username or team name where you're sending runs
  api_server: <OPTIONAL-WANDB-URL>    # Optional: In case your experiment tracking is not hosted at wandb servers

And that's it! Once you re-launch your training, ckpt_latest.pth (by default) will be downloaded to the checkpoints directory, and the training will resume from it just as if it was locally stored.

Evaluating Checkpoints Copy Chapter

Analogically to the previous section, we often want to evaluate a checkpoint seamlessly without being familiar with the training configuration.
For this reason, SG introduces two methods: Trainer.evaluate_checkpoint(...) and Trainer.evaluate_recipe(...) and play similar roles to the two previous ways of resuming experiments suggested in the last section:

Trainer.evaluate_checkpoint is used to evaluate a checkpoint resulting from one of your previous experiments, using the same parameters (dataset, valid_metrics,...) as used during the training of the experiment.
Trainer.evaluate_recipe is used to evaluate a checkpoint from SGs pre-trained model zoo or to evaluate a checkpoint with different parameters.
See both usages and documentation in the corresponding scripts evaluate_checkpoint and evaluate_recipe.

---

Documentation/Source/Configuration Files (documentation/source/configuration_files.md) Copy Section

Configuration Files and Recipes

defaults:
  - default_train_params
max_epochs: 250
lr_updates:
  _target_: numpy.arange
  start: 100
  stop: 250
  step: 50
lr_decay_factor: 0.1
lr_mode: step
lr_warmup_epochs: 0
initial_lr: 0.1
loss: CrossEntropyLoss
optimizer: SGD
criterion_params: {}
optimizer_params:
  weight_decay: 1e-4
  momentum: 0.9

train_dataset_params:
  root: /data/Imagenet/train
  transforms:
    - RandomResizedCropAndInterpolation:
        size: 224
        interpolation: default
    - RandomHorizontalFlip
    - ToTensor
    - Normalize:
        mean: [0.485, 0.456, 0.406] # mean for normalization
        std: [0.229, 0.224, 0.225]  # std  for normalization
val_dataset_params:
  root: /data/Imagenet/val
  transforms:
    - Resize:
        size: 256
    - CenterCrop:
        size: 224
    - ToTensor
    - Normalize:
        mean: [0.485, 0.456, 0.406] # mean for normalization
        std: [0.229, 0.224, 0.225]  # std  for normalization

python -m super_gradients.train_from_recipe --config-name=cifar10_resnet

The recipe you have just used is a configuration file containing everything SG needs to know in order to train
Resnet18 on Cifar10. The actual YAML file is located in src/super_gradients/recipes/cifar10_resnet.yaml.
In the same recipes library you can find many more configuration files defining different models, datasets,
and training hyper-parameters.

Hydra Copy Section

import hydra
from omegaconf import DictConfig
@hydra.main(config_path="recipes", version_base="1.2")
def main(cfg: DictConfig) -> None:
    print(cfg.experiment_name)

The @hydra.main decorator is looking for YAML files in the super_gradients.recipes according to the name of the configuration file provided
in the first arg of the command line.

In the experiment directory a .hydra subdirectory will be created. The configuration files related to this run will be saved by hydra to that subdirectory.

Two Hydra features worth mentioning are Command-Line Overrides
and YAML Composition.
We strongly recommend you to have a look at both of these pages.

Conclusion Copy Section

To handle data, SuperGradients takes use of two Pytorch primitives: torch.utils.data.Dataset - which is in charge of generating the samples and their corresponding labels,
and torch.utils.data.DataLoader - that wraps an iterable around the Dataset to enable easy access to the samples. In other words, torch.utils.data.Dataset defines how to load a single sample,
while torch.utils.data.DataLoader defines how to load batches of samples. For more information, see PyTorch documentation.

Datasets Copy Chapter

SuperGradients holds common public torch.utils.data.Dataset implementations for various tasks:

Classification:
Cifar10
Cifar100
ImageNetDataset

Object Detection:
COCODetectionDataset
DetectionDataset
PascalVOCDetectionDataset

Semantic Segmentation:
CoCoSegmentationDataSet
PascalAUG2012SegmentationDataSet
PascalVOC2012SegmentationDataSet
CityscapesDataset
SuperviselyPersonsDataset
PascalVOCAndAUGUnifiedDataset

Pose Estimation:
COCOKeypointsDataset

All of which can be imported from the super_gradients.training.datasets module. Note that some of the above implementations require following a few simple setup steps, which are all documented here

Creating a torch.utils.data.DataLoader from a dataset can be tricky, especially when defining some parameters on the fly. For example, in distributed training (i.e., DDP)
the torch.utils.data.DataLoader must be given a proper Sampler such that the dataset indices will be divided among the different processes.


Warning: Using the wrong sampler when defining a data loader to be used with DDP will lead the different processes to iterate over the same data samples giving little to no speedup over single GPU training!

from my_dataset import MyDataset
from super_gradients.training import dataloaders
import torchvision.transforms as T
from super_gradients.training import Trainer
from super_gradients.training.metrics import Accuracy
trainer = Trainer("my_experiment")
train_dataset = MyDataset(split="train", transforms=T.ToTensor())
valid_dataset = MyDataset(split="validation", transforms=T.ToTensor())
test_dataset = MyDataset(split="test", transforms=T.ToTensor())
train_dataloader = dataloaders.get(dataset=train_dataset, dataloader_params={"batch_size": 4})
valid_dataloader = dataloaders.get(dataset=valid_dataset, dataloader_params={"batch_size": 16})
test_dataloader = dataloaders.get(dataset=test_dataset, dataloader_params={"batch_size": 16})
model = ...
train_params = {...}
trainer.train(model=model, training_params=train_params, train_loader=train_dataloader, valid_loader=valid_dataloader)
trainer.test(model=trainer.net, test_loader=test_dataloader, test_metrics_list=[Accuracy()])

Note that dataloader_params will be unpacked in the torch.utils.data.DataLoader constructor after setting a proper sampler if one is not explicitly set.

DataLoaders Copy Chapter

As mentioned above, once instantiated, the torch.utils.data.DataLoader objects form batches.
Therefore- these are the objects being passed to Trainer.train(...):

...
trainer = Trainer("my_experiment")
train_dataloader = ...
valid_dataloader = ...
model = ...
train_params = {...}
trainer.train(model=model, training_params=train_params, train_loader=train_dataloader, valid_loader=valid_dataloader)

For your convenience, SuperGradients gives full access to all data loader objects used in our training recipes.
These are simply the torch.utils.data.DataLoader configured by the recipe's dataset_params:

cifar10_val
cifar10_train
cifar100_val
cifar100_train
coco2017_train
coco2017_val
coco2017_train_ssd_lite_mobilenet_v2
coco2017_val_ssd_lite_mobilenet_v2
imagenet_train
imagenet_val
imagenet_efficientnet_train
imagenet_efficientnet_val
imagenet_mobilenetv2_train
imagenet_mobilenetv2_val
imagenet_mobilenetv3_train
imagenet_mobilenetv3_val
imagenet_regnetY_train
imagenet_regnetY_val
imagenet_resnet50_train
imagenet_resnet50_val
imagenet_resnet50_kd_train
imagenet_resnet50_kd_val
imagenet_vit_base_train
imagenet_vit_base_val
tiny_imagenet_train
tiny_imagenet_val
pascal_aug_segmentation_train
pascal_aug_segmentation_val
pascal_voc_segmentation_train
pascal_voc_segmentation_val
supervisely_persons_train
supervisely_persons_val
pascal_voc_detection_train
pascal_voc_detection_val

These DataLoader can be imported from the super_gradients.training.dataloaders module.
Please note that these Dataset and DataLoader objects are already pre-defined with parameters required for specific training recipes. You can override these default parameters by passing two named arguments: dataset_params and dataloader_params(both of which are dictionaries), which will override the recipe settings. To learn which parameters you can override for each object, please refer to the YAML file with the same name.

For example, the code below will instantiate the data loader used for training in our imagenet_resnet50 recipe
(including all data augmentations and any other data-related setting which we defined for training Resnet50 on Imagenet)
but changing the batch size for our needs.
We can then, also with a one-liner, instantiate the validation dataloader and call train() as always:

from super_gradients.training.dataloaders import imagenet_resnet50_train, imagenet_resnet50_val
from super_gradients.training import Trainer
train_dataloader = imagenet_resnet50_train(dataloader_params={"batch_size": 4, "shuffle": True}, dataset_params={"root": "/my_data_dir/Imagenet/train"})
valid_dataloader = imagenet_resnet50_val(dataloader_params={"batch_size": 16}, dataset_params={"root": "/my_data_dir/Imagenet/val"})
...
trainer = Trainer("my_imagenet_training_experiment")
model = ...
train_params = {...}
trainer.train(model=model, training_params=train_params, train_loader=train_dataloader, valid_loader=valid_dataloader)

DataLoaders - Training with Configuration Files Copy Section

If you are still getting familiar with training with configuration files, follow this link.

Their names can reference any of the SG-predefined data loaders listed earlier.
For example, using the imagenet_resnet50_train and imagenet_resnet50_val:

dataset_params: ...
...
train_dataloader: imagenet_resnet50_train
val_dataloader: imagenet_resnet50_val
...

Now, on the structure of dataset_params:


train_dataset_params:
train_dataloader_params:
val_dataset_params:
val_dataloader_params:

train_dataloader: imagenet_resnet50_train
val_dataloader: imagenet_resnet50_val
dataset_params:
    train_dataset_params:
      root: /my_data_dir/Imagenet/train
    train_dataloader_params:
      batch_size: 4
      shuffle: True
    val_dataset_params:
      root: /my_data_dir/Imagenet/val
    val_dataloader_params:
      batch_size: 16

In addition to the parameters that are supported by the torch.utils.data.DataLoader class, SuperGradients also provide additional parameters:

* min_samples - When present, this parameter will guarantee that at least min_samples items will be processed in each epoch. It is useful when working with small datasets.
To use this option, simply add this parameter to the dataloader_params dictionary, and set it to the desired value:


train_dataloader: imagenet_resnet50_train
dataset_params:
    train_dataloader_params:
      batch_size: 4
      shuffle: True
      min_samples: 1024

On the technical side, when this parameter is se, SuperGradients will attach the RandomSampler to the DataLoader, and set it's num_samples parameter to min_samples.

Using Custom Datasets Copy Chapter

Suppose we already have our own torch.utils.data.Dataset class:


import torch
class MyCustomDataset(torch.utils.data.Dataset):
    def __init__(self, train: bool, image_size: int):
        ...
    def __getitem__(self, item):
        ...
        return inputs, targets # Or inputs, targets, additional_batch_items

#### A. __getitem__
You need to make sure that the __getitem__ method of your dataset complies with the following format:
- inputs = batch_items[0] : model input - The type might depend on the model you are using.
- targets = batch_items[1] : Target that will be used to compute loss/metrics - The type might depend on the function you are using.
- [OPTIONAL] additional_batch_items = batch_items[2] : Dict made of any additional item that you might want to use.

#### B. Train with your dataset
For coded training launch, we can instantiate it, then use it in the same way as the first code snippet to create
the data loaders and call train():

from my_dataset import MyCustomDataset
from super_gradients.training import dataloaders, Trainer
train_dataset = MyCustomDataset(train=True, image_size=64)
valid_dataset = MyCustomDataset(train=False, image_size=128)
train_dataloader = dataloaders.get(dataset=train_dataset, dataloader_params={"batch_size": 4, "shuffle": True})
valid_dataloader = dataloaders.get(dataset=valid_dataset, dataloader_params={"batch_size": 16})
trainer = Trainer("my_custom_dataset_training_experiment")
model = ...
train_params = {...}
trainer.train(model=model, training_params=train_params, train_loader=train_dataloader, valid_loader=valid_dataloader)

Using Custom Datasets - Training with Configuration Files Copy Section

When using configuration files, for example, training using train_from_recipe (or similar, when the underlying train method that is being called is Trainer.train_from_config(...)), In your my_dataset.py, register your dataset class by decorating the class with register_dataset:


import torch
from super_gradients.common.registry.registry import register_dataset
@register_dataset("my_custom_dataset")
class MyCustomDataset(torch.utils.data.Dataset):
    def __init__(self, train: bool, image_size: int):
        ...

dataset_params:
    train_dataset_params:
      train: True
      image_size: 64
    train_dataloader_params:
      dataset: my_custom_dataset
      batch_size: 4
      shuffle: True
    val_dataset_params:
      train: False
      image_size: 128
    val_dataloader_params:
      dataset: my_custom_dataset
      batch_size: 16

Last, in your my_train_from_recipe_script.py file, import the newly registered class (even though the class itself is unused, just to trigger the registry):



from omegaconf import DictConfig
  import hydra
  import pkg_resources
  from my_dataset import MyCustomDataset
  from super_gradients import Trainer, init_trainer
  
  
  @hydra.main(config_path=pkg_resources.resource_filename("super_gradients.recipes", ""), version_base="1.2")
  def main(cfg: DictConfig) -> None:
      Trainer.train_from_config(cfg)
  
  
  def run():
      init_trainer()
      main()
  
  
  if __name__ == "__main__":
      run()

Adding test datasets Copy Section

In addition to the train and validation datasets, you can also add a test dataset or multiple test datasets to your configuration file.
At the end of training, metrics from each test dataset will be computed and returned in final results.

#### Single test dataset

To add a single test dataset to recipe, add following properties to your configuration file:

test_dataloaders: <dataloader_name>
dataset_params:
  test_dataset_params:
    ...
  test_dataloader_params:
    ...


#### Multiple test datasets

In addition to the train and validation datasets, you can also add a test dataset or multiple test datasets to your configuration file.
This is how you can achieve this using YAML file:

#### Explicitly specifying all parameters

test_dataloaders:
  test_dataset_name_1: <dataloader_name>
  test_dataset_name_2: <dataloader_name>
dataset_params:
  test_dataset_params:
   test_dataset_name_1:
    ...
   test_dataset_name_2:
    ...
  test_dataloader_params:
   test_dataset_name_1:
    ...
   test_dataset_name_2:
    ...

#### Without dataloader names

A test_dataloaders property of the configuration file is optional and can be skipped.
You may want to use this option when you don't have a dataloader factory method registered.
In this case you have to specify a dataset class in corresponding dataloaders params.

dataset_params:
  test_dataset_params:
   test_dataset_name_1:
    ...
   test_dataset_name_2:
    ...
  test_dataloader_params:
   test_dataset_name_1:
     dataset: <dataset_class_name>
     ...
   test_dataset_name_2:
     dataset: <dataset_class_name>
     ...

#### Without dataloader params

A dataset_params.test_dataloader_params property is optional and can be skipped.
In this case dataset_params.val_dataloader_params will be used for instantiating test dataloaders.
Please note that if you don't use test_dataloaders and test_dataloader_params properties, a dataset_params.val_dataloader_params
must contain a dataset property specifying class name of the dataset to use.

dataset_params:
  test_dataset_params:
   test_dataset_name_1:
    ...
   test_dataset_name_2:
    ...

---

Documentation/Source/Dataloader Adapter (documentation/source/dataloader_adapter.md) Copy Section

Dataset Adapter

With diverse dataset structures available, ensuring compatibility with SuperGradients (SG) can be challenging. This is where the DataloaderAdapter plays a pivotal role. This tutorial takes you through the importance, implementation, and advantages of using the DataloaderAdapter with SG.

Why Dataset Adapter? Copy Section

Datasets come in a myriad of structures. However, SG requires data in a specific format.

For instance, consider the Object Detection Format:

Image format should be: (BS, H, W, C) i.e., channel last.
Targets should be in the format: (BS, 6), where 6 represents (sample_id, class_id, label, cx, cy, w, h).
The overhead of adjusting each dataset manually can be cumbersome. Enter DataloaderAdapter – designed to automatically understand your dataset structure and mold it for SG compatibility.

Why Do We Need the Dataset Adapter? Copy Section

While Datasets come in various structures and formats, SG expects data in a specific format to be able to run.

Example: Object Detection Format

To reduce this overhead, SuperGradients introduced the concept of DataloaderAdapter. Instead of requiring you to write all the transformations required to use SG, the DataloaderAdapter will infer anything possible directly from your data. Whenever something cannot be inferred with 100% confidence, you will be asked a question with all the required context for you to properly answer.

Let's see this in practice with an example. Let's start with SBDataset dataset

Exemple 1 - Segmentation Adapter on SBDataset Dataset

In this section, we'll walk through the process of preparing the SBDataset dataset for use in SuperGradients. We'll highlight the challenges and demonstrate how the Adapter can simplify the process.


1. Preparing the Dataset without Adapter

from torchvision.datasets import SBDataset
try:
  # There is a bug with torchvision.datasets.SBDataset that raises RuntimeError after downloading, so we just ignore it
  SBDataset(root="data", mode='segmentation', download=True)
except RuntimeError:
  pass

Downloading https://www2.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/semantic_contours/benchmark.tgz to data/benchmark.tgz


100%|██████████| 1419539633/1419539633 [00:32<00:00, 43301819.66it/s]


Extracting data/benchmark.tgz to data
Downloading https://www.cs.cornell.edu/~bharathh/ to data/train_noval.txt


20563it [00:00, 1012436.88it/s]

from torchvision.transforms import Compose, ToTensor, Resize, InterpolationMode

transforms = Compose([ToTensor(), Resize((512, 512), InterpolationMode.NEAREST)])
def sample_transform(image, mask):
  return transforms(image), transforms(mask)
train_set = SBDataset(root="data", mode='segmentation', download=False, transforms=sample_transform)

Now let's see what we get when instantiating a Dataloader

from torch.utils.data import DataLoader
train_loader = DataLoader(train_set, batch_size=20, shuffle=True)
_images, labels = next(iter(train_loader))
labels.unique()


tensor([0.0000, 0.0118, 0.0157, 0.0196, 0.0235, 0.0275, 0.0353, 0.0431, 0.0471,
0.0549, 0.0588, 0.0627, 0.0706, 0.0745, 0.0784])

As you can see, the labels are normalized (0-1). This is all right, but it is not the format expected by SuperGradients.

Let's now see how the Adapter helps.

2. Introducing Adapter

The Adapter helps us skip manual data preparations and dives right into creating a dataloader that SuperGradients expects.

from super_gradients.training.dataloaders.adapters import SegmentationDataloaderAdapterFactory
train_loader = SegmentationDataloaderAdapterFactory.from_dataset(dataset=train_set, batch_size=20, shuffle=True, config_path='local_cache.json')
_images, labels = next(iter(train_loader))
labels.unique()

[2023-10-29 15:25:36] INFO - data_config.py - Cache deactivated for SegmentationDataConfig.



--------------------------------------------------------------------------------
How many classes does your dataset include?
--------------------------------------------------------------------------------

Enter your response >>> 21
Great! You chose: 21

--------------------------------------------------------------------------------
Does your dataset provide a batch or a single sample?
--------------------------------------------------------------------------------
- Image shape: torch.Size([3, 512, 512])
- Mask shape: torch.Size([1, 512, 512])
Options:
[0] | Batch of Samples (e.g. torch Dataloader)
[1] | Single Sample (e.g. torch Dataset)

Your selection (Enter the corresponding number) >>> 1
Great! You chose: Single Sample (e.g. torch Dataset)

--------------------------------------------------------------------------------
In which format are your images loaded ?
--------------------------------------------------------------------------------

Options:
[0] | RGB
[1] | BGR
[2] | LAB
[3] | Other

Your selection (Enter the corresponding number) >>> 0
Great! You chose: RGB

tensor([ 0, 1, 2, 3, 4, 7, 8, 9, 12, 13, 15, 16, 18, 19, 20])

You can see that the mask is now encoded as int, which is the representation used in SuperGradients.

It's important to note that the dataset adapter also support different dataset format such as one hot, ensuring that the output (labels here) is in the right format to use within SuperGradients.

Example II - Detection Adapter on a Dictionary based Dataset Copy Chapter

Some datasets return a more complex data structure than the previous example.

For instance, the COCO dataset implementation from pytorch returns a list of dictionaries representing the labels.

Let's have a look:

Download the zip file

!wget https://deci-pretrained-models.s3.amazonaws.com/coco2017_small.zip
Unzip the downloaded file

!unzip coco2017_small.zip > /dev/null

--2023-10-29 15:27:31-- https://deci-pretrained-models.s3.amazonaws.com/coco2017_small.zip
Resolving deci-pretrained-models.s3.amazonaws.com (deci-pretrained-models.s3.amazonaws.com)... 54.231.134.129, 52.217.71.68, 52.217.138.65, ...
Connecting to deci-pretrained-models.s3.amazonaws.com (deci-pretrained-models.s3.amazonaws.com)|54.231.134.129|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 246116231 (235M) [application/zip]
Saving to: ‘coco2017_small.zip’

coco2017_small.zip 100%[===================>] 234.71M 38.7MB/s in 6.7s

2023-10-29 15:27:38 (34.9 MB/s) - ‘coco2017_small.zip’ saved [246116231/246116231]


from torchvision.datasets import CocoDetection
from torchvision.transforms import Compose, ToTensor, Resize, InterpolationMode
from torchvision.datasets import SBDataset

image_transform = Compose([ToTensor(), Resize((512, 512))])
train_set = CocoDetection(root='coco2017_small/images/train2017', annFile='coco2017_small/annotations/instances_train2017.json', transform=image_transform)
train_set = CocoDetection(root='coco2017_small/images/val2017', annFile='coco2017_small/annotations/instances_val2017.json', transform=image_transform)
image, targets = next(iter(train_set))

loading annotations into memory...
Done (t=0.10s)
creating index...
index created!
loading annotations into memory...
Done (t=0.05s)
creating index...
index created!

print(f"Number of targets: {len(targets)}, First target structure: {targets[0]}")

Observe the dataset output's nested dictionary structure? This complicates things for the Dataset Adapter as it's unsure which fields detail the bounding box.

To solve this, we utilize an extractor function.

#### The Extractor's Role

Simply put, the extractor translates your dataset's output into a format the Adapter understands. For our dataset, it will take the image and annotations, then return the bounding box data, including the label and coordinates.

Worried about bounding box format like xyxy_label or label_xywh? Don't be. The Adapter is designed to recognize them.

For further guidance on extractor functions, see the official documentation.

import torch
def coco_labels_extractor(sample) -> torch.Tensor:
    _, annotations = sample # annotations = [{"bbox": [1.08, 187.69, 611.59, 285.84], "category_id": 51}, ...]
    labels = []
    for annotation in annotations:
        class_id = annotation["category_id"]
        bbox = annotation["bbox"]
        labels.append((class_id, *bbox))
    return torch.Tensor(labels) # torch.Tensor([[51, 1.08, 187.69, 611.59, 285.84], ...])
coco_labels_extractor(sample=next(iter(train_set)))


tensor([[ 64.0000, 236.9800, 142.5100, 24.7000, 69.5000],
[ 72.0000, 7.0300, 167.7600, 149.3200, 94.8700],
[ 72.0000, 557.2100, 209.1900, 81.3500, 78.7300],
[ 62.0000, 358.9800, 218.0500, 56.0000, 102.8300],
[ 62.0000, 290.6900, 218.0000, 61.8300, 98.4800],
[ 62.0000, 413.2000, 223.0100, 30.1700, 81.3600],
[ 62.0000, 317.4000, 219.2400, 21.5800, 11.5900],
[ 1.0000, 412.8000, 157.6100, 53.0500, 138.0100],
[ 1.0000, 384.4300, 172.2100, 15.1200, 35.7400],
[ 78.0000, 512.2200, 205.7500, 14.7400, 15.9700],
[ 82.0000, 493.1000, 174.3400, 20.2900, 108.3100],
[ 84.0000, 604.7700, 305.8900, 14.3400, 45.7100],
[ 84.0000, 613.2400, 308.2400, 12.8800, 46.4400],
[ 85.0000, 447.7700, 121.1200, 13.9700, 21.8800],
[ 86.0000, 549.0600, 309.4300, 36.6800, 89.6700],
[ 86.0000, 350.7600, 208.8400, 11.3700, 22.5500],
[ 62.0000, 412.2500, 219.0200, 9.6300, 12.5200],
[ 86.0000, 241.2400, 194.9900, 14.2200, 17.6300],
[ 86.0000, 336.7900, 199.5000, 9.7300, 16.7300],
[ 67.0000, 321.2100, 231.2200, 125.5600, 88.9300]])

This output is all you need to get started. Now we can use the Dataloader Adapters!

from super_gradients.training.dataloaders.adapters import DetectionDataloaderAdapterFactory
from data_gradients.dataset_adapters.config.data_config import DetectionDataConfig

adapter_config = DetectionDataConfig(labels_extractor=coco_labels_extractor, cache_path="coco_adapter_cache.json")
train_loader = DetectionDataloaderAdapterFactory.from_dataset(
    dataset=train_set,
    config=adapter_config,
    batch_size=5,
    drop_last=True,
)
val_loader = DetectionDataloaderAdapterFactory.from_dataset(
    dataset=train_set,
    config=adapter_config,
    batch_size=5,
    drop_last=True,
)

/usr/local/lib/python3.10/dist-packages/torchvision/transforms/functional.py:1603: UserWarning: The default value of the antialias parameter of all the resizing transforms (Resize(), RandomResizedCrop(), etc.) will change from None to True in v0.17, in order to be consistent across the PIL and Tensor backends. To suppress this warning, directly pass antialias=True (recommended, future default), antialias=None (current default, which means False for Tensors and True for PIL), or antialias=False (only works on Tensors - PIL will still use antialiasing). This also applies if you are using the inference transforms from the models weights: update the call to weights.transforms(antialias=True).
warnings.warn(
[2023-10-29 15:27:41] INFO - data_config.py - Cache deactivated for DetectionDataConfig.
[2023-10-29 15:27:41] INFO - detection_adapter_collate_fn.py - You are using Detection Adapter. Please note that it was designed specifically for YOLONAS, YOLOX and PPYOLOE.


Number of targets: 20, First target structure: {'segmentation': [[240.86, 211.31, 240.16, 197.19, 236.98, 192.26, 237.34, 187.67, 245.8, 188.02, 243.33, 176.02, 250.39, 186.96, 251.8, 166.85, 255.33, 142.51, 253.21, 190.49, 261.68, 183.08, 258.86, 191.2, 260.98, 206.37, 254.63, 199.66, 252.51, 201.78, 251.8, 212.01]], 'area': 531.8071000000001, 'iscrowd': 0, 'image_id': 139, 'bbox': [236.98, 142.51, 24.7, 69.5], 'category_id': 64, 'id': 26547}

--------------------------------------------------------------------------------
How many classes does your dataset include?
--------------------------------------------------------------------------------

Enter your response >>> 80
Great! You chose: 80

--------------------------------------------------------------------------------
In which format are your images loaded ?
--------------------------------------------------------------------------------

Options:
[0] | RGB
[1] | BGR
[2] | LAB
[3] | Other

Your selection (Enter the corresponding number) >>> 0
Great! You chose: RGB

--------------------------------------------------------------------------------
Which comes first in your annotations, the class id or the bounding box?
--------------------------------------------------------------------------------
Here's a sample of how your labels look like:
Each line corresponds to a bounding box.
tensor([[ 64.0000, 236.9800, 142.5100, 24.7000, 69.5000],
[ 72.0000, 7.0300, 167.7600, 149.3200, 94.8700],
[ 72.0000, 557.2100, 209.1900, 81.3500, 78.7300],
[ 62.0000, 358.9800, 218.0500, 56.0000, 102.8300]])
Options:
[0] | Label comes first (e.g. [class_id, x1, y1, x2, y2])
[1] | Bounding box comes first (e.g. [x1, y1, x2, y2, class_id])

Your selection (Enter the corresponding number) >>> 0
Great! You chose: Label comes first (e.g. [class_id, x1, y1, x2, y2])

--------------------------------------------------------------------------------
What is the bounding box format?
--------------------------------------------------------------------------------
Here's a sample of how your labels look like:
Each line corresponds to a bounding box.
tensor([[ 64.0000, 236.9800, 142.5100, 24.7000, 69.5000],
[ 72.0000, 7.0300, 167.7600, 149.3200, 94.8700],
[ 72.0000, 557.2100, 209.1900, 81.3500, 78.7300],
[ 62.0000, 358.9800, 218.0500, 56.0000, 102.8300]])
Options:
[0] | xyxy: x-left, y-top, x-right, y-bottom (Pascal-VOC format)
[1] | xywh: x-left, y-top, width, height (COCO format)
[2] | cxcywh: x-center, y-center, width, height (YOLO format)

Your selection (Enter the corresponding number) >>> 1


[2023-10-29 15:28:40] INFO - detection_adapter_collate_fn.py - You are using Detection Adapter. Please note that it was designed specifically for YOLONAS, YOLOX and PPYOLOE.

_image, targets = next(iter(train_loader))

targets.shape # [N, 6] format with 6 representing (sample_id, class_id, cx, cy, w, h)


torch.Size([22, 6])

targets[:3]


tensor([[ 0.0000, 64.0000, 249.3300, 177.2600, 24.7000, 69.5000],
[ 0.0000, 72.0000, 81.6900, 215.1950, 149.3200, 94.8700],
[ 0.0000, 72.0000, 597.8850, 248.5550, 81.3500, 78.7300]])

III. Use your Adapted Dataloader to train a model

Now that we have an adapter for a detection dataset, let's use it to launch a training of YoloNAS on it!

This is of course for the sake of the example, since YoloNAS was originally trained using the SuperGradients implementation of COCO Dataset. You can replace the COCO dataset with any of your dataset.

from omegaconf import OmegaConf
from hydra.utils import instantiate
from super_gradients import Trainer
from super_gradients.training import models
from super_gradients.common.object_names import Models
from super_gradients.training import training_hyperparams
from super_gradients.common.environment.cfg_utils import load_recipe
trainer = Trainer(experiment_name="yolonas_training_with_adapter", ckpt_root_dir="../../scripts/")
model = models.get(model_name=Models.YOLO_NAS_S, num_classes=adapter_config.n_classes, pretrained_weights="coco")
yolonas_recipe = load_recipe(config_name="coco2017_yolo_nas_s",
                             overrides=[f"arch_params.num_classes={adapter_config.n_classes}",
                                        "training_hyperparams.max_epochs=1",
                                        "training_hyperparams.mixed_precision=False"])
yolonas_recipe = OmegaConf.to_container(instantiate(yolonas_recipe))
training_params = yolonas_recipe['training_hyperparams']
trainer.train(model=model, training_params=training_params, train_loader=train_loader, valid_loader=val_loader)

[2023-10-29 15:29:42] INFO - checkpoint_utils.py - License Notification: YOLO-NAS pre-trained weights are subjected to the specific license terms and conditions detailed in
https://github.com/Deci-AI/super-gradients/blob/master/LICENSE.YOLONAS.md
By downloading the pre-trained weight files you agree to comply with these terms.
Downloading: "https://sghub.deci.ai/models/yolo_nas_s_coco.pth" to /root/.cache/torch/hub/checkpoints/yolo_nas_s_coco.pth
100%|██████████| 73.1M/73.1M [00:00<00:00, 81.0MB/s]
[2023-10-29 15:29:43] INFO - checkpoint_utils.py - Successfully loaded pretrained weights for architecture yolo_nas_s
[2023-10-29 15:29:44] INFO - sg_trainer.py - Starting a new run with run_id=RUN_20231029_152944_310569
[2023-10-29 15:29:44] INFO - sg_trainer.py - Checkpoints directory: ./yolonas_training_with_adapter/RUN_20231029_152944_310569
[2023-10-29 15:29:44] INFO - sg_trainer.py - Using EMA with params {'decay': 0.9997, 'decay_type': 'threshold', 'beta': 15}


The console stream is now moved to ./yolonas_training_with_adapter/RUN_20231029_152944_310569/console_Oct29_15_29_44.txt


[2023-10-29 15:29:45] WARNING - callbacks.py - Number of warmup steps (1000) is greater than number of steps in epoch (100). Warmup steps will be capped to number of steps in epoch to avoid interfering with any pre-epoch LR schedulers.
/usr/local/lib/python3.10/dist-packages/torchvision/transforms/functional.py:1603: UserWarning: The default value of the antialias parameter of all the resizing transforms (Resize(), RandomResizedCrop(), etc.) will change from None to True in v0.17, in order to be consistent across the PIL and Tensor backends. To suppress this warning, directly pass antialias=True (recommended, future default), antialias=None (current default, which means False for Tensors and True for PIL), or antialias=False (only works on Tensors - PIL will still use antialiasing). This also applies if you are using the inference transforms from the models weights: update the call to weights.transforms(antialias=True).
warnings.warn(
/usr/local/lib/python3.10/dist-packages/super_gradients/training/utils/collate_fn/detection_collate_fn.py:29: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
images_batch = [torch.tensor(img) for img in images_batch]
/usr/local/lib/python3.10/dist-packages/super_gradients/training/utils/collate_fn/detection_collate_fn.py:43: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
labels_batch = [torch.tensor(labels) for labels in labels_batch]
[2023-10-29 15:29:45] INFO - sg_trainer_utils.py - TRAINING PARAMETERS:
- Mode: Single GPU
- Number of GPUs: 0 (0 available on the machine)
- Full dataset size: 500 (len(train_set))
- Batch size per GPU: 5 (batch_size)
- Batch Accumulate: 1 (batch_accumulate)
- Total batch size: 5 (num_gpus * batch_size)
- Effective Batch size: 5 (num_gpus batch_size batch_accumulate)
- Iterations per epoch: 100 (len(train_loader))
- Gradient updates per epoch: 100 (len(train_loader) / batch_accumulate)

[2023-10-29 15:29:45] INFO - sg_trainer.py - Started training for 1 epochs (0/0)

Train epoch 0: 100%|██████████| 100/100 [20:18<00:00, 12.18s/it, PPYoloELoss/loss=4, PPYoloELoss/loss_cls=1.72, PPYoloELoss/loss_dfl=2.2, PPYoloELoss/loss_iou=0.472, gpu_mem=0]
Validating: 100%|██████████| 100/100 [06:34<00:00, 3.94s/it]
[2023-10-29 15:56:43] INFO - base_sg_logger.py - Checkpoint saved in ./yolonas_training_with_adapter/RUN_20231029_152944_310569/ckpt_best.pth
[2023-10-29 15:56:43] INFO - sg_trainer.py - Best checkpoint overriden: validation [email protected]:0.95: 0.0005365016404539347


===========================================================
SUMMARY OF EPOCH 0
├── Train
│ ├── Ppyoloeloss/loss_cls = 1.7168
│ ├── Ppyoloeloss/loss_iou = 0.4717
│ ├── Ppyoloeloss/loss_dfl = 2.2035
│ └── Ppyoloeloss/loss = 3.9977
└── Validation
├── Ppyoloeloss/loss_cls = 2.4528
├── Ppyoloeloss/loss_iou = 0.5016
├── Ppyoloeloss/loss_dfl = 2.2003
├── Ppyoloeloss/loss = 4.807
├── [email protected]:0.95 = 0.0052
├── [email protected]:0.95 = 0.007
├── [email protected]:0.95 = 0.0005
└── [email protected]:0.95 = 0.0007

===========================================================


[2023-10-29 15:56:45] INFO - sg_trainer.py - RUNNING ADDITIONAL TEST ON THE AVERAGED MODEL...
Validating epoch 1: 100%|██████████| 100/100 [06:33<00:00, 3.93s/it]

IV. Dig deeper into the Adapter

By default, any parameter that could not be confidently infered will trigger a question.

But you have the possibility to set these parameters in advance through the config object. In the previous example we had to set labels_extractor explicitly. Now let's set all the parameters

from super_gradients.training.dataloaders.adapters import DetectionDataloaderAdapterFactory
from data_gradients.dataset_adapters.config.data_config import DetectionDataConfig
from data_gradients.utils.data_classes.image_channels import ImageChannels
class_names = [category['name'] for category in train_set.coco.loadCats(train_set.coco.getCatIds())]
adapter_config = DetectionDataConfig(
    labels_extractor=coco_labels_extractor,
    is_label_first=True,
    class_names=class_names,
    image_channels=ImageChannels.from_str("RGB"),
    xyxy_converter='xywh',
    cache_path="coco_adapter_cache_with_default.json"
)

This can now be used and you don't need to answer any question

train_loader = DetectionDataloaderAdapterFactory.from_dataset(
    dataset=train_set,
    config=adapter_config,
    batch_size=5,
    drop_last=True,
)
val_loader = DetectionDataloaderAdapterFactory.from_dataset(
    dataset=train_set,
    config=adapter_config,
    batch_size=5,
    drop_last=True,
)
_image, targets = next(iter(train_loader))
print(targets.shape) # [N, 6] format with 6 representing (sample_id, class_id, cx, cy, w, h)

[2023-10-29 16:15:09] INFO - detection_adapter_collate_fn.py - You are using Detection Adapter. Please note that it was designed specifically for YOLONAS, YOLOX and PPYOLOE.
/usr/local/lib/python3.10/dist-packages/torchvision/transforms/functional.py:1603: UserWarning: The default value of the antialias parameter of all the resizing transforms (Resize(), RandomResizedCrop(), etc.) will change from None to True in v0.17, in order to be consistent across the PIL and Tensor backends. To suppress this warning, directly pass antialias=True (recommended, future default), antialias=None (current default, which means False for Tensors and True for PIL), or antialias=False (only works on Tensors - PIL will still use antialiasing). This also applies if you are using the inference transforms from the models weights: update the call to weights.transforms(antialias=True).
warnings.warn(
[2023-10-29 16:15:09] INFO - detection_adapter_collate_fn.py - You are using Detection Adapter. Please note that it was designed specifically for YOLONAS, YOLOX and PPYOLOE.
/usr/local/lib/python3.10/dist-packages/super_gradients/training/utils/collate_fn/detection_collate_fn.py:29: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
images_batch = [torch.tensor(img) for img in images_batch]
/usr/local/lib/python3.10/dist-packages/super_gradients/training/utils/collate_fn/detection_collate_fn.py:43: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
labels_batch = [torch.tensor(labels) for labels in labels_batch]

Load from existing cache Copy Section

You can use the cache of an adapter you already used in the past. This will allow you skip the questions that were already asked in the previous run.

The new config will load the answer from questions asked in the previous run.

adapter_config = DetectionDataConfig(
    labels_extractor=coco_labels_extractor,
    cache_path="coco_adapter_cache_with_default.json" # Name of the previous cache
)
train_loader = DetectionDataloaderAdapterFactory.from_dataset(
    dataset=train_set,
    config=adapter_config,
    batch_size=5,
    drop_last=True,
)
val_loader = DetectionDataloaderAdapterFactory.from_dataset(
    dataset=train_set,
    config=adapter_config,
    batch_size=5,
    drop_last=True,
)
_image, targets = next(iter(train_loader))

targets.shape # [N, 6] format with 6 representing (sample_id, class_id, cx, cy, w, h)


torch.Size([22, 6])

As you can see, no question was asked and we still get the targets adapted into the SuperGradients format.

---

Documentation/Source/DetectionOutputAdapter (documentation/source/DetectionOutputAdapter.md) Copy Section

DetectionOutputAdapter

The DetectionOutputAdapter is a class that converts the output of a detection model into a user-appropriate format.
For instance, it can be used to convert the format of bounding boxes from CYXHW to XYXY, or to change the layout of the elements
in the output tensor from [X1, Y1, X2, Y2, Confidence, Class] to [Class, Confidence, X1, Y1, X2, Y2].

Features Copy Chapter

* Easy rearrangement of the elements in the output tensor
* Easy conversion of the bounding box format
* Support of JIT Tracing & Scripting
* Support of ONNX export

Usage Copy Chapter

We start by introducing the concept of a format. A format represents a specific layout of the elements in the output tensor.
Currently, there is only one type of formats supported - ConcatenatedTensorFormat which represents a layout where all predictions concatenated into a single tensor.
Additional formats can be added in the future (Like DictionaryOfTensorsFormat).

ConcatenatedTensorFormat requires that input is a tensor and has the following shape:
* Tensor of shape [N, Elements] - N is the number of predictions, Elements is the concatenated vector of attributes per box.
* Tensor of shape [B, N, Elements] - B is the batch dimension, N and Elements as above.

To instantiate the DetectionOutputAdapter we have to describe the input and output formats for our predictions:

Let's imagine model emits predictions in the following format:

[N, 10] (cx, cy, w, h, class, confidence, attributes..)

example_input = [
    #      cx          cy        w          h     class, confidence,   attribute a, attribute b, attribute c, attribute d
    [0.465625,  0.5625,    0.13125,   0.125,          0,      0.968,         0.350,       0.643,       0.640,       0.453],
    [0.103125,  0.1671875, 0.10625,   0.134375,       1,      0.897,         0.765,       0.654,       0.324,       0.816],
    [0.078125,  0.078125,  0.15625,   0.15625,        2.,     0.423,         0.792,       0.203,       0.653,       0.777],
    ...
]

The corresponding format definition would look like this:

from super_gradients.training.datasets.data_formats import ConcatenatedTensorFormat, BoundingBoxesTensorSliceItem, TensorSliceItem, NormalizedCXCYWHCoordinateFormat
input_format = ConcatenatedTensorFormat(
    layout=(
        BoundingBoxesTensorSliceItem(name="bboxes", format=NormalizedCXCYWHCoordinateFormat()),
        TensorSliceItem(name="class", length=1),
        TensorSliceItem(name="confidence", length=1),
        TensorSliceItem(name="attributes", length=4),
    )
)

For sake of demonstration, let's assume that we want to convert the output to the following format:

[N, 10] (class, attributes, x1, y1, x2, y2)

[
    # class, attribute a, attribute b, attribute c, attribute d,     x1,   y1,   x2,    y2
    [     0,       0.350,       0.643,       0.640,       0.453,    256,  320,  340,   400],
    [     1,       0.765,       0.654,       0.324,       0.816,     32,   64,  100,   150],
    [     2,       0.792,       0.203,       0.653,       0.777,      0,    0,  100,   100],
    ...
]

* The class and attributes are the same as in the input format but comes first
* The format of bounding boxes is changed from NormalizedCXCYWHCoordinateFormat to XYXYCoordinateFormat
* The confidence is removed from the output

The corresponding format definition would look like this:

from super_gradients.training.datasets.data_formats import ConcatenatedTensorFormat, BoundingBoxesTensorSliceItem, TensorSliceItem, XYXYCoordinateFormat
output_format = ConcatenatedTensorFormat(
    layout=(
        TensorSliceItem(name="class", length=1),
        TensorSliceItem(name="attributes", length=4),
        BoundingBoxesTensorSliceItem(name="bboxes", format=XYXYCoordinateFormat()),
    )
)

Now we can construct the DetectionOutputAdapter and attach it to the model:

from super_gradients.training.datasets.data_formats import DetectionOutputAdapter
output_adapter = DetectionOutputAdapter(input_format, output_format, image_shape=(640,640))
model = nn.Sequential(
    create_model(),
    create_nms(),
    output_adapter
)

To test how the output adapter transforms dummy input one can easily run it alone:

output = output_adapter(torch.from_numpy(example_input)).numpy()
print(output)
Prints:

[
    # class,   attribute a, attribute b, attribute c, attribute d,     x1,   y1,   x2,    y2
    [     0,         0.350,       0.643,       0.640,       0.453,    256,  320,  340,   400], 
    [     1,         0.765,       0.654,       0.324,       0.816,     32,   64,  100,   150], 
    [     2,         0.792,       0.203,       0.653,       0.777,      0,    0,  100,   100]
]

Not supported features Copy Chapter

Currently DetectionOutputAdapter does not support the following features:

* argmax operation over a slice of confidences for [C] classes (Useful to compute argmax(class confidences))
Multiplication of two slices (Useful to compute confidence class)

---

Documentation/Source/Device (documentation/source/device.md) Copy Section

Training Modes

SuperGradients allows users to train models on different modes:
1. CPU
2. single GPU - (CUDA)
3. multiple GPUs' - Data Parallel (DP)
4. multiple GPUs' - Distributed Data Parallel (DDP)

1. CPU Copy Chapter

How to use it: If you don't have any CUDA device available, your training will automatically be run on CPU.
Otherwise, the default device will be CUDA, but you can still easily set it to CPU using setup_device as follow:


from super_gradients import Trainer
from super_gradients.training.utils.distributed_training_utils import setup_device
setup_device(device='cpu')
Unchanged

trainer = Trainer(...)
trainer.train(...)

2. CUDA Copy Chapter

How to use it: All you need to do is to call a magic function setup_device before instantiating the Trainer.


from super_gradients import Trainer
from super_gradients.training.utils.distributed_training_utils import setup_device
Launch DP on 4 GPUs'

setup_device(multi_gpu='DP', num_gpus=4)
Unchanged

trainer = Trainer(...)
trainer.train(...)

How to use it: All you need to do is to call a magic function setup_device before instantiating the Trainer.


from super_gradients import Trainer
from super_gradients.training.utils.distributed_training_utils import setup_device
Launch DDP on 4 GPUs'

setup_device(num_gpus=4) # Equivalent to: setup_device(multi_gpu='DDP', num_gpus=4)
Unchanged

trainer = Trainer(...)
trainer.train(...)

In the following example, the print_hello function will print Hello world only once, when it would be printed 4 times without the decorator...


from super_gradients.training.utils.distributed_training_utils import setup_device
from super_gradients.common.environment.ddp_utils import multi_process_safe
setup_device(num_gpus=4)
@multi_process_safe # Try with and without this decorator
def print_hello():
    print('Hello world')
print_hello()


#### B. DDP requires specific Metric implementation!
As explained, multiple processes are used to train a model with DDP, each on its own GPU.
This means that the metrics must be computed and aggregated across all the processes, and it requires the metric to be implemented using states.

States are attributes to be reduced. They are defined using the built-in method add_state() and enable broadcasting
of the states among the different ranks when calling the compute() method.

An example of a state would be the number of correct predictions, which will be summed across the different processes, and broadcasted to all of
them before computing the metric value. You can see an example below.

Feel free to check torchmetrics documentation for more information on how to implement your own metric.

Example
In the following example, we start with a custom metric implemented to run on a single device:


import torch
from torchmetrics import Metric

class Top5Accuracy(Metric):
    def __init__(self):
        super().__init__()
        self.correct = torch.tensor(0.)
        self.total = torch.tensor(0.)
    def update(self, preds: torch.Tensor, target: torch.Tensor):
        batch_size = target.size(0)
        # Get the top k predictions
        _, pred = preds.topk(5, 1, True, True)
        pred = pred.t()
        # Count the number of correct predictions only for the highest 5
        correct = pred.eq(target.view(1, -1).expand_as(pred))
        correct5 = correct[:5].reshape(-1).float().sum(0)
        self.correct += correct5.cpu()
        self.total += batch_size
    def compute(self):
        return self.correct.float() / self.total

import torch
import torchmetrics
class DDPTop1Accuracy(torchmetrics.Metric):
    def __init__(self, dist_sync_on_step=False):
        super().__init__(dist_sync_on_step=dist_sync_on_step)
        self.add_state("correct", default=torch.tensor(0.), dist_reduce_fx="sum")   # Set correct to be a state
        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")      # Set total to be a state
    def update(self, preds: torch.Tensor, target: torch.Tensor):
        batch_size = target.size(0)
        # Get the top k predictions
        _, pred = preds.topk(5, 1, True, True)
        pred = pred.t()
        # Count the number of correct predictions only for the highest 5
        correct = pred.eq(target.view(1, -1).expand_as(pred))
        correct5 = correct[:5].reshape(-1).float().sum(0)
        self.correct += correct5
        self.total += batch_size
    def compute(self):
        return self.correct.float() / self.total

python -m super_gradients.train_from_recipe -m --config-name=coco2017_yolo_nas_s training_hyperparams.initial_lr=1e-3,5e-3,1e-4

training_recipe.yaml

default:
   - ...
...
Simply add this to run DDP on 4 nodes.

gpu_mode: DDP
num_gpus: 4

ema_param = ema_param  decay + param  (1 - decay)


To enable use of EMA is SuperGradients one should add following parameters to the training_params:

from super_gradients import Trainer
trainer = Trainer(...)
trainer.train(
    training_params={"ema": True, "ema_params": {"decay": 0.9999, "decay_type": "constant"}, ...}, 
    ...
)

The decay is a hyperparameter that controls the speed of the EMA update. It's value must be in (0,1) range.
Larger values of decay will result in slower EMA model update.

It is usually beneficial to have smaller decay values at the start of training and increase it as the training progresses.
In SuperGradients we support several types of changing decay value over time:

- constant: "ema_params": {"decay": 0.9999, "decay_type": "constant"}
- threshold: "ema_params": {"decay": 0.9999, "decay_type": "threshold"}
- exp: "ema_params": {"decay": 0.9999, "decay_type": "exp", "beta": 15}

Adding your own decay schedule Copy Chapter

It is possible to bring your own decay schedule in SuperGradients. By subclassing from IDecayFunction one can implement a custom
function:

from super_gradients.training.utils.ema_decay_schedules import IDecayFunction, EMA_DECAY_FUNCTIONS
class LinearDecay(IDecayFunction):
    def __init__(self, kwargs):
        pass
    def __call__(self, decay: float, step: int, total_steps: int) -> float:
      """
      Compute EMA for specific training step following linear scaling rule [0..decay)
        :param decay: The maximum decay value.
        :param step: Current training step. The unit-range training percentage can be obtained by step / total_steps.
        :param total_steps:  Total number of training steps.
        :return: Computed decay value for a given step.
      """
      training_progress = step / total_steps
      return decay * training_progress
EMA_DECAY_FUNCTIONS["linear"] = LinearDecay

How EMA weights are saved Copy Chapter

When EMA is enabled, saved checkpoints will contain additional ema_net attribute.
Weights for EMA model are saved under ema_net.
A regular (non-averaged) model weights are saved as net key in checkpoint as usual.

When instantiating a model via models.get, a function will check whether ema_net is present in the checkpoint.
In case ema_net is in checkpoint, the model will be initialized using EMA weights; otherwise a model will load initialized from regular weights saved in net.

Knowledge Distillation Copy Chapter

EMA is also supported in knowledge distillation. To enable it one should add following parameters to the training_params similar to the above example:

from super_gradients import KDTrainer
trainer = KDTrainer(...)
trainer.train(
    training_params={"ema": True, "ema_params": {"decay": 0.9999, "decay_type": "constant"}, ...}, 
    ...
)

---

Documentation/Source/Example Classification (documentation/source/Example_Classification.md) Copy Section

Training a Classification Model and Transfer Learning

In this example we will use SuperGradients to train from scratch a ResNet18 model on the CIFAR10 image classification
dataset. We will also fine-tune the same model via transfer learning with weights pre-trained on the ImageNet dataset.

Quick installation Copy Chapter

For this example, the only necessary package is super-gradients. Installing super-gradients will also
install all dependencies required to run the code in this example.

pip install super-gradients

1. Experiment setup Copy Chapter

First, we will initialize the Trainer. It handles:
- Model training
- Evaluating test data
- Making predictions
- Saving and managing checkpoints


To initialize it, you need:

- Experiment Name: A unique identifier for your training experiment.
- Checkpoint Root Directory (ckpt_root_dir): The directory where checkpoints, logs, and tensorboards are saved. While optional, if unspecified, it assumes the presence of a 'checkpoints' directory in your project's root.

from super_gradients import Trainer
experiment_name = "resnet18_cifar10_example"
CHECKPOINT_DIR = '/path/to/checkpoints/root/dir'
trainer = Trainer(experiment_name=experiment_name, ckpt_root_dir=CHECKPOINT_DIR)

2. Understanding the Checkpoint Structure Copy Section

Checkpoints are crucial for progressive training, debugging, and model deployment. SuperGradients organizes them in a structured manner. Here's what the directory hierarchy looks like under your specified ckpt_root_dir:

<ckpt_root_dir>
│
├── <experiment_name>
│   │
│   ├─── <run_dir>
│   │     ├─ ckpt_best.pth                   # Best performance during validation
│   │     ├─ ckpt_latest.pth                 # End of the most recent epoch
│   │     ├─ average_model.pth               # Averaged over specified epochs
│   │     ├─ ckpt_epoch_*.pth                # Checkpoints from specific epochs (like epoch 10, 15, etc.)
│   │     ├─ events.out.tfevents.*           # Tensorflow run artifacts
│   │     └─ log_<timestamp>.txt             # Trainer logs of the specific run
│   │
│   └─── <other_run_dir>
│        └─ ...
│
└─── <other_experiment_name>
    │
    ├─── <run_dir>
    │     └─ ...
    │
    └─── <another_run_dir>
          └─ ...

In this structure:

- ckpt_best.pth: Saved whenever there's an improvement in the specified validation metric.
- ckpt_latest.pth: Updated at the end of every epoch.
- average_model.pth: Averaged checkpoint, created if average_best_models parameter is set to True.

For more information, check out the dedicated page.

2. Dataset and dataloaders Copy Chapter

The dataset used in this example is the CIFAR10 image classification dataset.

SuperGradients provides a pool of standard datasets and dataloaders readily available for quick and easy usage.
SuperGradients also downloads the datasets when necessary, and gracefully handles the creation of the dataloaders,
with a pre-made training recipe specifically tailored for the dataset and model architecture.

<ins>Note:</ins> The SuperGradients trainer is compatible with PyTorch dataloaders and dataset objects. While it
is outside the scope of this example, it is worth remembering that custom dataloaders and datasets
can be employed when necessary.

2.A. Default dataloader from SuperGradients Copy Section

As can be seen in the code snippet below, creating the training and validation dataloaders using SuperGradients'
default implementation is as easy as writing two lines of code:

from super_gradients.training import dataloaders
train_dataloader = dataloaders.get(name="cifar10_train", dataset_params={}, dataloader_params={"num_workers": 2})
valid_dataloader = dataloaders.get(name="cifar10_val", dataset_params={}, dataloader_params={"num_workers": 2})

Here, we call the get() function twice, for the training and validation dataloaders. The function's parameters are:

* name - a string representing the name of the desired dataloader, out of a variety of different pre-made dataloaders
provided by SuperGradients. In this example, we use the pre-made CIFAR10 training and validation dataloaders.
* dataset_params - a dictionary of dataset-related parameters. Used to override the default parameters defined in the
training recipe. Later in this example we will show how this can be used to change the transforms
applied to the images.
* dataloader_params - a dictionary of dataloader-related parameters. Used to override the default parameters defined
in the training recipe. Here, as an example, we set the number of workers to 2.
* dataset - a torch.utils.data.Dataset object. Used when employing a custom dataset implementation. This parameter
cannot be passed together with the name or dataset_params parameters.

We can always print the parameter values of the dataloader and its related dataset:

import pprint
print("Dataloader parameters:")
pprint.pprint(train_dataloader.dataloader_params)
print("Dataset parameters:")
pprint.pprint(train_dataloader.dataset.dataset_params)

Expected output:

Dataloader parameters:
{
    "batch_size": 256,
    "drop_last": False,
    "num_workers": 2,
    "pin_memory": True,
    "shuffle": True
}
Dataset parameters:
{
    "download": True,
    "root": "./data/cifar10",
    "target_transform": None,
    "train": True,
    "transforms": [
        {"RandomCrop": {"size": 32, "padding": 4}},
        "RandomHorizontalFlip",
        "ToTensor",
        {"Normalize": {"mean": [0.4914, 0.4822, 0.4465], "std": [0.2023, 0.1994, 0.201]}},
    ]
}

When the get() function is called as above, SuperGradients will attempt to download the CIFAR10 dataset for us. We
can expect to see an output as follows:

<img src="images/classification_example_dataset_download.png" width="750">

After the dataloaders are defined, we can iterate them to extract batches of images and their corresponding labels.
This is useful for several purposes, such as visualization, verifying tensor shapes, and more. For example, visualization:

from matplotlib import pyplot as plt
def show(images, labels, classes, rows=6, columns=5):
  fig = plt.figure(figsize=(10, 10))
  for i in range(1, columns * rows + 1):
      fig.add_subplot(rows, columns, i)
      plt.imshow(images[i-1].permute(1, 2, 0).clamp(0, 1))
      plt.xticks([])
      plt.yticks([])
      plt.title(f"{classes[labels[i-1]]}")
      
  plt.show()
      
images_train, labels_train = next(iter(train_dataloader))
show(images_train, labels_train, classes=train_dataloader.dataset.classes)

Output:

<img src="images/classification_example_data_vis.png" width="500">

As can be seen, the images are normalized. The normalization process is defined, among other things, as part of the
default training recipe SuperGradients uses for the CIFAR10 dataset. As we will see in the following section,
SuperGradients makes it a trivial task to override all, or part, of the different dataset and dataloader
parameters, allowing for control over the flexibility vs. ease-of-use tradeoff.

For completion of this section, let's print the tensors' shapes:

print(f'Training image tensor shape: {images_train.shape}')
print(f'Training labels tensor shape: {labels_train.shape}')

output:

Training image tensor shape: torch.Size([256, 3, 32, 32])
Training labels tensor shape: torch.Size([256])

As we can see, the default batch size for the training dataloader is 256.

2.B. Override parameters in dataset and dataloaders creation Copy Section

To showcase the flexibility SuperGradients allows in customizing the different trainer
components, we will override the list of transforms that are used in the dataset. To define a list of
transformations to apply, we will use torchvision's transforms. This also serves to show
the seamless integration SuperGradients allows with different PyTorch components. For the sake of
visualization, the only transform we apply is ToTensor(), which simply converts the input
images into PyTorch tensors.

from torchvision import transforms as T
transforms_list = [T.ToTensor()]
vis_dataloader = dataloaders.get("cifar10_train",
                                 dataset_params={"transforms": transforms_list},
                                 dataloader_params={"num_workers": 2})
images, labels = next(iter(vis_dataloader))
show(images, labels, classes=train_dataloader.dataset.classes)

Notice that the only difference in the dataloader's definition is that here, the dataset_params parameter
is passed as a dictionary defining the parameters to override. The result of running the above code:

<img src="images/classification_example_transforms_data_vis.png" width="500">

The effect of changing the transforms can be seen in the images - now, without normalization, the objects are more
observable in the images.

3. Architecture definition Copy Chapter

In this example, we train the model with the ResNet18 architecture.

SuperGradients provides out-of-the-box implementations of many architectures for classification tasks. With just one
line of code we can define a model with the chosen architecture. A list of all available architectures can be found
here.

from super_gradients.training import models
from super_gradients.common.object_names import Models
model = models.get(model_name=Models.RESNET18, num_classes=10)

Notice that, similar to obtaining a pre-defined dataloader, here we use super_gradients.training.models's get()
function. In the above code, two parameters are passed to the function:

* model_name - A string defining the model's architecture name, out of the list of architectures SuperGradients
provides.
* num_classes - An integer representing the number of classes the model should learn to predict. Affects the
architecture's structure.

Some additional parameters the get() function supports:

* arch_params - A dictionary used to override the default architecture parameters, such as the number of residual
blocks.
* checkpoint_path - A string defining the path to an external checkpoint to be loaded. Can be absolute or relative.
If provided, will automatically attempt to load the checkpoint.
* pretrained_weights - A string defining the name of a dataset on which the model was pre-trained on, for
fine-tuning and transfer learning. The pretrained_weights and checkpoint_path parameters
are mutually exclusive.

For more available parameters, refer to the function's docstring.

In this example, we have defined the model with one of SuperGradient's readily available architectures. As was
already noted in previous sections, SuperGradients is highly compatible with PyTorch. Defining the model's architecture
is not an exception - we can seamlessly use a custom architecture, i.e. a torch.nn.Module object, for maximum
flexibility, although it is out of the scope of this example.

4. Training setup Copy Chapter

We have defined the trainer, datasets, dataloaders, and model architecture. Before we can start training, we need to
define the training parameters. As with the other parameters, SuperGradients provides training parameters optimized for
this use-case. For more recommended training parameters you can have a look at our recipes
here.

Obtaining the training parameters is as easy as writing a single line of code:

from super_gradients.training import training_hyperparams
training_params =  training_hyperparams.get(config_name="training_hyperparams/cifar10_resnet_train_params")

We notice the repeatability in the code usage - to obtain the training parameters, we again call
the get() function. This function accepts two parameters:

* config_name - A string defining the .yaml config filename in the recipes' directory.
* overriding_params - An optional parameter, a dictionary used to override the loaded training parameters.

We can print the training parameters to see the different options:

pprint.pprint("Training parameters:") 
pprint.pprint(training_params)

Output (Training parameters):

/ Detailed source-code truncated for AI context efficiency. /

As can be seen in the above output, there are numerous options to modify the training parameters
to affect the training process.

It is also possible to change training parameters after obtaining them, for example:

training_params["max_epochs"] = 15
training_params["sg_logger_params"]["launch_tensorboard"] = True

5. Training, checkpointing, and transfer learning Copy Chapter

5.A. Training the model Copy Section

We are all set to start training our model. Simply plug in the model, training and validation dataloaders,
and training parameters into the trainer's train() function:

trainer.train(model=model, 
              training_params=training_params, 
              train_loader=train_dataloader,
              valid_loader=valid_dataloader)

The training progress will be printed to the screen:

[2023-02-01 20:57:27] INFO - sg_trainer_utils.py - TRAINING PARAMETERS:
    - Mode:                         Single GPU
    - Number of GPUs:               1          (4 available on the machine)
    - Dataset size:                 50000      (len(train_set))
    - Batch size per GPU:           256        (batch_size)
    - Batch Accumulate:             1          (batch_accumulate)
    - Total batch size:             256        (num_gpus * batch_size)
    - Effective Batch size:         256        (num_gpus  batch_size  batch_accumulate)
    - Iterations per epoch:         195        (len(train_set) / total_batch_size)
    - Gradient updates per epoch:   195        (len(train_set) / effective_batch_size)
[2023-02-01 20:57:27] INFO - sg_trainer.py - Started training for 15 epochs (0/14)
Train epoch 0: 100%|██████████| 196/196 [00:18<00:00, 10.51it/s, Accuracy=0.262, CrossEntropyLoss=2.37, Top5=0.787, gpu_mem=0.371]
Validation epoch 0: 100%|██████████| 20/20 [00:03<00:00,  6.12it/s]
===========================================================
SUMMARY OF EPOCH 0
├── Training
│   ├── Accuracy = 0.262
│   ├── CrossEntropyLoss = 2.3702
│   └── Top5 = 0.787
└── Validation
    ├── Accuracy = 0.3459
    ├── CrossEntropyLoss = 1.8811
    └── Top5 = 0.871
===========================================================

At the beginning of the training, a summary of the training parameters is printed, where we can see
the training mode (CPU/single GPU/distributed training), the number of GPUs used, the training dataset
size, and more. The progress of each epoch's training and validation is displayed, along with the tracked metrics
(defined as part of the training recipe): accuracy, loss value, top5 error, and GPU memory consumption.
At the end of each epoch, a summary of the training and validation metrics is displayed, and in later epochs,
a comparison with the previous epochs is provided:

===========================================================
SUMMARY OF EPOCH 15
├── Training
│   ├── Accuracy = 0.7594
│   │   ├── Best until now = 0.7458 (↗ 0.0136)
│   │   └── Epoch N-1      = 0.7458 (↗ 0.0136)
│   ├── CrossEntropyLoss = 0.686
│   │   ├── Best until now = 0.7187 (↘ -0.0327)
│   │   └── Epoch N-1      = 0.7187 (↘ -0.0327)
│   └── Top5 = 0.9867
│       ├── Best until now = 0.9849 (↗ 0.0019)
│       └── Epoch N-1      = 0.9849 (↗ 0.0019)
└── Validation
    ├── Accuracy = 0.7425
    │   ├── Best until now = 0.746  (↘ -0.0035)
    │   └── Epoch N-1      = 0.7306 (↗ 0.0119)
    ├── CrossEntropyLoss = 0.7331
    │   ├── Best until now = 0.7315 (↗ 0.0016)
    │   └── Epoch N-1      = 0.8048 (↘ -0.0717)
    └── Top5 = 0.9831
        ├── Best until now = 0.9838 (↘ -0.0007)
        └── Epoch N-1      = 0.9818 (↗ 0.0013)
===========================================================

At the end of each epoch, the different logs and checkpoints are saved in the path defined by ckpt_root_dir and
experiment_name.

In the epoch summary shown above, we can see that the validation accuracy is 73%, which is
not very high. To get better insights as to what is happening, we turn to the tensorboard logs.

5.B. Tensorboard logs Copy Section

To view the experiment's tensorboard logs, type the following command in the terminal from the
experiment's path:

tensorboard --logdir='.'

(Alternatively, run the command from anywhere with the experiment's full path).

SuperGradients logs many useful metrics to tensorboard, including CPU and GPU usage, learning rate
scheduling, training and validation losses and other metrics, and many more. For the purpose of this example,
let us examine the training and validation loss:

<img src="images/classification_example_train_loss_initial.png" width="500">

<img src="images/classification_example_valid_loss_initial.png" width="500">

As can be seen in the graphs, the training (and validation) loss did not converge before training ended. This means
that training the model for additional epochs will probably improve its performance. Earlier, when modifying the
training parameters, we set max_epochs = 15. Let us continue training the model for an additional 10 epochs.

5.C. Continue training from a checkpoint Copy Section

To continue training from a checkpoint, we utilize the models.get() function's checkpoint_path parameter.
The provided checkpoint path includes the checkpoint file we wish to load. In this example, since we want
to continue from the last checkpoint, we will load the ckpt_latest.pth checkpoint. Additionally, we want to let
the trainer know that we are continuing training and not starting from the first epoch. This is done
by setting the resume training parameter to True. Finally, we set the new max_epochs training
parameter, and train the model once more.

import os
model = models.get(model_name=Models.RESNET18, 
                   num_classes=10,
                   checkpoint_path=os.path.join(CHECKPOINT_DIR, experiment_name, 'ckpt_latest.pth'))
training_params["resume"] = True
training_params["max_epochs"] = 25
trainer.train(model=model,
              training_params=training_params,
              train_loader=train_dataloader,
              valid_loader=valid_dataloader)

We can see that the training continues for 10 epochs, resuming from epoch 15:

[2023-02-01 21:21:16] INFO - sg_trainer_utils.py - TRAINING PARAMETERS:
    - Mode:                         Single GPU
    - Number of GPUs:               1          (4 available on the machine)
    - Dataset size:                 50000      (len(train_set))
    - Batch size per GPU:           256        (batch_size)
    - Batch Accumulate:             1          (batch_accumulate)
    - Total batch size:             256        (num_gpus * batch_size)
    - Effective Batch size:         256        (num_gpus  batch_size  batch_accumulate)
    - Iterations per epoch:         195        (len(train_set) / total_batch_size)
    - Gradient updates per epoch:   195        (len(train_set) / effective_batch_size)
[2023-02-01 21:21:16] INFO - sg_trainer.py - Started training for 10 epochs (15/24)
Train epoch 15: 100%|██████████| 196/196 [00:18<00:00, 10.52it/s, Accuracy=0.764, CrossEntropyLoss=0.668, Top5=0.987, gpu_mem=0.422]
Validation epoch 15: 100%|██████████| 20/20 [00:03<00:00,  6.10it/s]
===========================================================
SUMMARY OF EPOCH 15
├── Training
│   ├── Accuracy = 0.7644
│   ├── CrossEntropyLoss = 0.6684
│   └── Top5 = 0.9865
└── Validation
    ├── Accuracy = 0.7539
    ├── CrossEntropyLoss = 0.7271
    └── Top5 = 0.9841
===========================================================

Finally, the model stops training after completing 25 epochs:

===========================================================
SUMMARY OF EPOCH 25
├── Training
│   ├── Accuracy = 0.8177
│   │   ├── Best until now = 0.8147 (↗ 0.003)
│   │   └── Epoch N-1      = 0.8147 (↗ 0.003)
│   ├── CrossEntropyLoss = 0.5211
│   │   ├── Best until now = 0.5281 (↘ -0.007)
│   │   └── Epoch N-1      = 0.5281 (↘ -0.007)
│   └── Top5 = 0.9921
│       ├── Best until now = 0.9919 (↗ 0.0002)
│       └── Epoch N-1      = 0.9919 (↗ 0.0002)
└── Validation
    ├── Accuracy = 0.8201
    │   ├── Best until now = 0.7873 (↗ 0.0328)
    │   └── Epoch N-1      = 0.7534 (↗ 0.0667)
    ├── CrossEntropyLoss = 0.525
    │   ├── Best until now = 0.6145 (↘ -0.0895)
    │   └── Epoch N-1      = 0.7517 (↘ -0.2266)
    └── Top5 = 0.9907
        ├── Best until now = 0.9883 (↗ 0.0024)
        └── Epoch N-1      = 0.983  (↗ 0.0077)
===========================================================

We can see that the validation accuracy is now 82%. Much better.

5.C. Transfer learning Copy Section

So far, we trained a model from scratch. More formally, the model's weights were randomly initialized. For easy tasks
and large datasets this is usually sufficient. In other cases, especially when not a lot of data is available for
training, we would like to take advantage of knowledge gained from other sources. This is called transfer learning,
and it has many forms and variations. In this example we will take a look at the simplest form of transfer learning:
fine-tuning a model initialized with pre-trained weights. Specifically, we will initialize our model
with weights pre-trained on the ImageNet dataset.

SuperGradients provides a variety of pre-trained weights readily available for fine-tuning different models. To
initialize our model with pre-trained weights provided by SuperGradients, only a small change to the existing code
is needed:

model = models.get(model_name=Models.RESNET18, num_classes=10, pretrained_weights="imagenet")

In the above code, we provided the pretrained_weights parameter to the models.get() function. This parameter accepts
a string representing the name of the dataset that the weights were pre-trained on. Note that this
parameter is mutually exclusive with the checkpoint_path parameter.

The rest of the training pipeline is the same as above. For comparison with the previous model, we will train this model
for 25 epochs also. Let us look how the model performed:

===========================================================
SUMMARY OF EPOCH 25
├── Training
│   ├── Accuracy = 0.8242
│   │   ├── Best until now = 0.8267 (↘ -0.0025)
│   │   └── Epoch N-1      = 0.8267 (↘ -0.0025)
│   ├── CrossEntropyLoss = 0.5035
│   │   ├── Best until now = 0.4998 (↗ 0.0037)
│   │   └── Epoch N-1      = 0.4998 (↗ 0.0037)
│   └── Top5 = 0.9924
│       ├── Best until now = 0.9917 (↗ 0.0007)
│       └── Epoch N-1      = 0.9917 (↗ 0.0007)
└── Validation
    ├── Accuracy = 0.8377
    │   ├── Best until now = 0.8062 (↗ 0.0315)
    │   └── Epoch N-1      = 0.806  (↗ 0.0317)
    ├── CrossEntropyLoss = 0.4834
    │   ├── Best until now = 0.5731 (↘ -0.0897)
    │   └── Epoch N-1      = 0.5785 (↘ -0.0952)
    └── Top5 = 0.9924
        ├── Best until now = 0.9903 (↗ 0.0021)
        └── Epoch N-1      = 0.9903 (↗ 0.0021)
===========================================================

As we can see, the validation accuracy improved by 1.7% compared to the randomly initialized model.
To achieve a greater improvement with pre-trained weights, sometimes careful tuning of the training hyperparameters
is required.

6. Predictions with the trained model Copy Chapter

Now that we have a trained model with reasonable performance, we can use it to make predictions on new data.
First, let's import some packages:

from PIL import Image
import torch
import numpy as np
import requests

Next, we load the model with the trained weights, and put it into evaluation mode. Notice that now we load the
ckpt_best.pth checkpoint.

model = models.get(model_name=Models.RESNET18, 
                   num_classes=10,
                   checkpoint_path=os.path.join(CHECKPOINT_DIR, experiment_name, 'ckpt_best.pth'))
model.eval()

We want to test the model on an image of one the classes the model was trained on. As an example, let us test how the
model handles an image of a frog, taken from the Aquarium of the Pacific website.
The loaded image must undergo the same transformations as the training images for the model to work well:

url = "https://www.aquariumofpacific.org/images/exhibits/Magnificent_Tree_Frog_900.jpg"
image = np.array(Image.open(requests.get(url, stream=True).raw))
transforms = T.Compose([
    T.ToTensor(),
    T.Normalize(mean=(0.4914, 0.4822, 0.4465), std=(0.2023, 0.1994, 0.2010)),
    T.Resize((32, 32))
])
input_tensor = transforms(image).unsqueeze(0).to(next(model.parameters()).device)

Next, to obtain the model's predictions we simply run the following line of code:

predictions = model(input_tensor)

Let's see what the model predicted:

plt.xlabel(train_dataloader.dataset.classes[torch.argmax(predictions)])
plt.imshow(image)
plt.show()

<img src="images/classification_example_frog_prediction.png" width="500">

As we can see, the model correctly predicted that the input image is an image of a frog.

7. Complete code Copy Chapter

For completion of this example, we provide a complete working code for training, continuing training from
a saved checkpoint, and predicting with the trained model. Simply change the CHECKPOINT_DIR variable and run the
script:

/ Detailed source-code truncated for AI context efficiency. /

---

Documentation/Source/Example Training An External Model (documentation/source/Example_Training-an-external-model.md) Copy Section

Training an external model

In this example we will use SuperGradients to train a deep learning segmentation model to extract human portraits from
images, i.e., to remove the background from the image.

We will show how SuperGradients allows seamless integration of
an external model, dataset, loss function, and metric into the training pipeline.

Quick installation Copy Chapter

For this example, the only necessary package is super-gradients. Installing super-gradients will also
install all dependencies required to run the code in this example.

pip install super-gradients

1. Dataset Copy Chapter

The dataset we will use in this example is the
AISegment dataset, available to download for free from
Kaggle.
The dataset contains 34,427 RGB images of human portraits and their corresponding soft masks.
The provided images are center-cropped to a unified shape of 600x800.

1.A. Data preparation Copy Section

The original structure of the data on the disk is as follows:

data
└───aisegment-matting
│   └───matting
│       │   └───1803290511
│       │   │       └───matting_00000000
│       │   │       │         1803290511-00000459.png
│       │   │       │         1803290511-00000458.png
│       │   │       │         ..
│       │   └───1803290444
│       │       ..
│   └───clip_img
│       │   └───1803290511
│       │   │       └───clip_00000000
│       │   │       │         1803290511-00000459.jpg
│       │   │       │         1803290511-00000458.jpg
│       │   │       │         ..
│       │   └───1803290444
│       │       ..

In the structure shown above, the clip_img directory contains the portrait images, divided into many
sub-directories. The matting directory is structured similar to the clip_img directory, containing
the corresponding masks. This structure is not particularly convenient for data loading, therefore
we will first rearrange the data with the following script, which only requires changing the data_path and
out_path variables:

/ Detailed source-code truncated for AI context efficiency. /

The original data may have images without their corresponding masks. Therefore, the above code first finds all existing
image-mask pairs. It then splits the data into train/validation/test sets, where the size of each of the validation and
test sets is 10% of the total number of samples. The samples in each of the sets are then copied to the output path,
structured as follows:

data
    └───train
    │       1803290511-00000029.jpg
    │       1803290511-00000029.png
    │       ..
    └───val
    │       1803290444-00000292.jpg
    │       1803290444-00000292.png
    │       ..
    └───test
    │       1803290443-00000294.jpg
    │       1803290443-00000294.png
    │       ..

Each of the train/val/test directories contain all image (.jpg) files and their corresponding mask (.png) files.

1.B. PyTorch Dataset Copy Section

In some cases, we may want to have full control over the process of loading and pre-processing the training data.
SuperGradients is fully compatible with PyTorch data loaders, which allows for seamless integration of custom dataset
implementations for maximum flexibility.

We will first present the complete Dataset class implementation, and then break it down to fully understand what is
going on.

/ Detailed source-code truncated for AI context efficiency. /

First, we can see that our dataset class inherits from the torch.utils.data.Dataset class. To initialize the dataset
object, two parameters must be provided:

* data_path - the full path to the data's root directory
* split - a string indicating whether this is the 'train', 'val', or 'test' split of the data

Additional parameters include input_height and input_width, the fixed height and width, respectively, for resizing
the input images and masks. These are set to 256 by default. At the end of the initialization function, the image and
mask path pairs are determined according to the split.

Next, let's see what happens when we retrieve an item from the dataset via the __getitem__() function.
First, an image and its corresponding mask are loaded according to the idx parameter:

image = Image.open(self.path_pairs[idx][0]).convert('RGB')
mask = Image.open(self.path_pairs[idx][1]).split()[-1]

Notice that we take only the mask's last channel. The mask's original color format is RGBA, and the alpha channel should be
used as the soft mask for segmentation according to the
dataset's description.

Next, we apply transformations to the image and mask according to the split:

if self.split == 'train':
    seg_transforms = torch_transforms.Compose([
        SegRandomFlip(prob=0.5),
        SegColorJitter(brightness=0.5, contrast=0.5, saturation=0.5),
        SegResize(h=self.input_height, w=self.input_width)
    ])
else:
    seg_transforms = SegResize(h=self.input_height, w=self.input_width)
transformed_pair = seg_transforms({"image": image, "mask": mask})
image, mask = transformed_pair['image'], transformed_pair['mask']

On training images and masks we apply data augmentation: a random horizontal flip with probability 0.5, and
random color jitter which randomly changes the image's brightness, contrast, and saturation. Both the training and the
other splits' images and masks undergo resizing according to input_height and input_width.

The transforms in the above code are all SuperGradients transforms, which are built upon PyTorch's
torchvision.transforms. This allows for maximum compatibility with PyTorch components. For example, we can see that
the transforms are sequentially composed using torchvision's Compose class. Notice also that the transform names are
all prefixed with Seg - meaning that they are specifically designed for segmentation data. These transforms take care
to apply the same transformation to the image and the mask when required, or only apply the transformation to the image
otherwise. For example, if SegRandomFlip() flips the image, the mask will be flipped as well. SegColorJitter() only
transforms the image. SegResize() resizes both the image and the mask.

Next, we convert the image and the mask to PyTorch Tensors, and normalize the image using the NORMALIZATION_MEANS and
NORMALIZATION_STDS:

image_transform = torch_transforms.Compose([
    torch_transforms.ToTensor(),
    torch_transforms.Normalize(self.NORMALIZATION_MEANS, self.NORMALIZATION_STDS)
])
mask_transform = torch_transforms.ToTensor()
image_tensor, mask_tensor = image_transform(image), mask_transform(mask)

These transforms are applied regardless of the current dataset split. Notice that here we use torchvision transforms.
This serves to show the high degree of flexibility SuperGradients allows, as its transforms are based on PyTorch's
torchvision transforms which may be used interchangeably.

1.C. Data visualization Copy Section

To conclude this section, let's visualize some images and their masks to test our AISegmentDataset implementation.
First, we instantiate two dataset objects, for the training and validation splits, and extract the first sample from
each:

import matplotlib.pyplot as plt
data_path = '/path/to/arranged/data/dir'
train_dataset = AISegmentDataset(data_path=data_path, split='train')
val_dataset = AISegmentDataset(data_path=data_path, split='val')
train_image, train_mask = train_dataset[0]
val_image, val_mask = val_dataset[0]

Let's first visualize the validation image and mask:

figure = plt.figure()
figure.add_subplot(1, 2, 1)
plt.title("Image")
plt.axis("off")
plt.imshow(val_image.permute(1, 2, 0))
figure.add_subplot(1, 2, 2)
plt.title("Mask")
plt.axis("off")
plt.imshow(val_mask.squeeze(), cmap='gray')
plt.show()

<img src="images/external_model_example_val_vis.png" width="750">

And the training image and mask:

figure = plt.figure()
figure.add_subplot(1, 2, 1)
plt.title("Image")
plt.axis("off")
plt.imshow(train_image.permute(1, 2, 0))
figure.add_subplot(1, 2, 2)
plt.title("Mask")
plt.axis("off")
plt.imshow(train_mask.squeeze(), cmap='gray')
plt.show()

<img src="images/external_model_example_train_vis.png" width="750">

We can see the effect of the color jitter transform on the image.

2. Model architecture Copy Chapter

For this example we will employ the U-Net architecture. U-Net and its variants are a
popular choice for many image segmentation tasks. It is a fully-convolutional architecture,
consisting of an expanding and a contracting path with skip connections between the encoder
and decoder blocks. In this example we will demonstrate how we can easily integrate an external PyTorch model
as part of SuperGradients' training pipeline. To this end, we will use a
U-Net implementation loaded directly from
PyTorch Hub:

model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch', 'unet', in_channels=3, out_channels=1, 
                       init_features=32, pretrained=False)

Since our model's inputs are RGB images, we set in_channels=3. This is a binary segmentation task, therefore
out_channels=1. The init_features parameter determines the number of kernels in the first block's convolution
layers. The number of kernels is doubled in each consecutive encoder block. We also set pretrained=False since
in this case we do not want to use pre-trained weights.

Let's check our model's type:

print(type(model))

<class 'unet.UNet'>

SuperGradients is compatible with models of type torch.nn.Module. Just to be sure, let's verify that our
model is of the correct type:

print(type(model).__bases__)

(<class 'torch.nn.modules.module.Module'>,)

Finally, let's print the model to see its components:

print(model)

/ Detailed source-code truncated for AI context efficiency. /

3. Loss function Copy Chapter

While a popular choice for binary segmentation loss function is the binary cross-entropy (BCE) loss, often it is
the intersection-over-union (IoU), or the Jaccard Index, that serves
as a measure of success. In this example, we will use a combination of the BCE and IoU as the loss function. While
SuperGradients provides, among many other losses, an implementation of the combined BCE and Dice loss, which could be
used for our purposes as well, we will show how we can define our own user-defined loss function and train our model
with it using SuperGradients.

Similar to using an external model, the custom loss function's class must inherit from torch.nn.Module. The
forward() function's first parameter needs to be the predictions tensor and the second parameter needs to be the
target tensor.

import torch
import torch.nn as nn

class CustomIoU(torch.nn.Module):
    def __init__(self):
        super(CustomIoU, self).__init__()
    def forward(self, preds, target):
        intersection = torch.sum(target * preds)
        union = torch.sum(target) + torch.sum(preds) - intersection + 1e-5
        iou = intersection / union
        return iou

class CustomSegLoss(torch.nn.Module):
    def __init__(self, bce_weight=1, iou_weight=1):
        super(CustomSegLoss, self).__init__()
        self.bce_weight = bce_weight
        self.iou_weight = iou_weight
        self.bce_loss = nn.BCELoss()
        self.iou_func = CustomIoU()
    def forward(self, preds, target):
        bce_loss = self.bce_loss(preds, target)
        iou_loss = 1.0 - self.iou_func(torch.gt(preds, 0.5).long(), torch.gt(target, 0.5).long())
        return self.bce_weightbce_loss + self.iou_weightiou_loss

Notice that here the BCE loss term is obtained simply by using PyTorch's BCELoss(). To compute the IoU score, we
have implemented an auxiliary class CustomIoU, which implements a naive, differentiable IoU function. Note that
in binary segmentation tasks, we are usually interested only in the foreground IoU. Therefore, CustomIoU disregards
the background IoU. To compute the IoU, the pixel values in both images should be binary, i.e., 0's and 1's.
Since in the model's forward function a sigmoid function is already applied to the output, we only need to binarize
both the predictions and the target tensors with a threshold of 0.5 (remember, we are using soft masks). When
measuring segmentation performance, higher IoU is better. Since IoU score is a number in [0, 1], we simply compute
IoU loss = 1 - IoU to make it a valid loss function for gradient-descent optimization.

The overall loss is a weighted sum of the BCE and the IoU loss terms, with weights bce_weight and iou_weight,
respectively.

With just a few lines of code we have defined our own custom loss function. Although here we could have used a similar
loss function provided by SuperGradients, in many other cases a more complex and specific loss function is desired.
This example serves to show how we may define any loss function to fit our needs and seamlessly integrate it with the
SuperGradients training pipeline.

4. Custom IoU metric Copy Chapter

Since our measure of success is the IoU, we would like to tell SuperGradients to log and track this metric during
training. This metric would also be used to determine the best model at the end of every epoch for checkpointing.
We note that SuperGradients provides many built-in metrics, including variants of the IoU. However, in this example
we aim to show the ease at which we can incorporate external metrics into our pipeline. SuperGradients supports any
metric of type torchmetrics.Metric.

The metric we will use is torchmetrics' JaccardIndex. However, recall that we use soft masks.
JaccardIndex requires the target
tensor's elements to be integers. Also, as noted in the previous section, since this is a binary segmentation
task, we are interested in the foreground IoU. Therefore, we will need to modify the metric a bit. For more details
about implementing a custom metric using torchmetrics, see
here.

class SoftIoU(JaccardIndex):
    def __init__(self, kwargs):
        super().__init__(reduction='none', kwargs)
    def update(self, preds: torch.Tensor, target: torch.Tensor):
        target = torch.gt(target, 0.5).long()
        super().update(preds, target)
    def compute(self):
        return super().compute()[1]

We have defined our SoftIoU class which inherits from JaccardIndex. The only modifications we introduced to
JaccardIndex are:

1. The target tensor is binarized with a threshold of 0.5: target = torch.gt(target, 0.5).long()
2. To get the IoU of the foreground alone, we set reduction='none' in the __init__ function. This means that
instead of computing the mean of the background and foreground IoUs, both values are returned, and in the compute()
function we only take the second element, which corresponds to the foreground.

Our custom metric is now ready to use with our training pipeline.

5. Experiment configuration Copy Chapter

Trainer Copy Section

- Experiment Name: A unique identifier for your training experiment.
- Checkpoint Root Directory (ckpt_root_dir): The directory where checkpoints, logs, and tensorboards are saved. While optional, if unspecified, it assumes the presence of a 'checkpoints' directory in your project's root.

from super_gradients import Trainer
experiment_name = "aisegment_example"
CHECKPOINT_DIR = '/path/to/checkpoints/root/dir'
trainer = Trainer(experiment_name=experiment_name, ckpt_root_dir=CHECKPOINT_DIR)

Understanding the Checkpoint Structure Copy Section

Checkpoints are crucial for progressive training, debugging, and model deployment. SuperGradients organizes them in a structured manner. Here's what the directory hierarchy looks like under your specified ckpt_root_dir:

<ckpt_root_dir>
│
├── <experiment_name>
│   │
│   ├─── <run_dir>
│   │     ├─ ckpt_best.pth                   # Best performance during validation
│   │     ├─ ckpt_latest.pth                 # End of the most recent epoch
│   │     ├─ average_model.pth               # Averaged over specified epochs
│   │     ├─ ckpt_epoch_*.pth                # Checkpoints from specific epochs (like epoch 10, 15, etc.)
│   │     ├─ events.out.tfevents.*           # Tensorflow run artifacts
│   │     └─ log_<timestamp>.txt             # Trainer logs of the specific run
│   │
│   └─── <other_run_dir>
│        └─ ...
│
└─── <other_experiment_name>
    │
    ├─── <run_dir>
    │     └─ ...
    │
    └─── <another_run_dir>
          └─ ...

In this structure:

- ckpt_best.pth: Saved whenever there's an improvement in the specified validation metric.
- ckpt_latest.pth: Updated at the end of every epoch.
- average_model.pth: Averaged checkpoint, created if average_best_models parameter is set to True.

For more information, check out the dedicated page.

Dataloaders Copy Section

Next, we initialize the PyTorch dataloaders for our datasets:

from torch.utils.data import DataLoader
train_dataloader = DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=2)
val_dataloader = DataLoader(val_dataset, batch_size=16, shuffle=False, num_workers=2)

Training Hyperparameters Copy Section

And lastly, we need to define the training hyperparameters:

train_params = {
    "max_epochs": 100,
    "lr_mode": "CosineLRScheduler",
    "initial_lr": 0.001,
    "optimizer": "Adam",
    "loss": CustomSegLoss(),
    "metric_to_watch": "SoftIoU",
    "greater_metric_to_watch_is_better": True,
    "train_metrics_list": [SoftIoU(num_classes=2)],
    "valid_metrics_list": [SoftIoU(num_classes=2)]
}

Notice that the training hyperparameters must be defined as a dictionary with the hyperparameter names as keys.
The dictionary defines the hyperparameters that we want to override. All other hyperparameters retain their default
values defined by SuperGradients. The list of all training hyperparameters and their default value
can be found here.

The metric_to_watch hyperparameter defines the metric used to determine the best model at the end of every epoch.
We simply set it as a string representing the name of our custom metric SoftIoU. Greater IoU is better, therefore we
set greater_metric_to_watch_is_better=True. The IoU metric will also be logged and tracked during training and
validation, as determined by train_metrics_list and valid_metrics_list. The loss hyperparameter tells
SuperGradients which loss function to use during training. To use one of the many loss functions provided by
SuperGradients, we set this hyperparameter as a string representing the loss function's name. However, here we want
to use our custom loss function. Therefore, we simply set the hyperparameter as a CustomSegLoss() object.

The above code shows the simplicity of integrating external, user-defined components into the SuperGradients training
pipeline. We simply plugged instantiations of our custom loss and metric into the hyperparameters dictionary,
and we are ready to go.

6. Training Copy Chapter

6.A. Training the model Copy Section

We are all set to start training our model. Simply plug in the model, training and validation dataloaders,
and training parameters into the trainer's train() function:

trainer.train(model=model, 
              training_params=train_params, 
              train_loader=train_dataloader,
              valid_loader=val_dataloader)

The training progress will be printed to the screen:

[2023-02-06 11:44:35] INFO - sg_trainer.py - Started training for 100 epochs (0/99)
Train epoch 0: 100%|██████████| 3443/3443 [19:16<00:00,  2.98it/s, CustomSegLoss=0.24, SoftIoU=0.9, gpu_mem=1.81]
Validation epoch 0: 100%|██████████| 431/431 [00:42<00:00, 10.23it/s]
===========================================================
SUMMARY OF EPOCH 0
├── Training
│   ├── Customsegloss = 0.2466
│   └── Softiou = 0.9
└── Validation
    ├── Customsegloss = 0.1367
    └── Softiou = 0.9483
===========================================================

The progress of each epoch's training and validation is displayed, along with the values of our custom metric and
loss function, and GPU memory consumption. At the end of each epoch, a summary of the training and validation metrics
is displayed, and in later epochs, a comparison with the previous epochs is provided:

===========================================================
SUMMARY OF EPOCH 5
├── Training
│   ├── Customsegloss = 0.0915
│   │   ├── Best until now = 0.0945 (^[[32m↘ -0.003^[[0m)
│   │   └── Epoch N-1      = 0.0945 (^[[32m↘ -0.003^[[0m)
│   └── Softiou = 0.9651
│       ├── Best until now = 0.9    (↗ 0.0651^[[0m)
│       └── Epoch N-1      = 0.9639 (↗ 0.0011^[[0m)
└── Validation
    ├── Customsegloss = 0.0789
    │   ├── Best until now = 0.0924 (^[[32m↘ -0.0135^[[0m)
    │   └── Epoch N-1      = 0.0924 (^[[32m↘ -0.0135^[[0m)
    └── Softiou = 0.9702
        ├── Best until now = 0.9483 (↗ 0.0219^[[0m)
        └── Epoch N-1      = 0.9657 (↗ 0.0045^[[0m)
===========================================================

At the end of each epoch, the different logs and checkpoints are saved in the path defined by ckpt_root_dir and
experiment_name. Let's see how we can use Tensorboard to track training process.

6.B. Tensorboard logs Copy Section

To view the experiment's tensorboard logs, type the following command in the terminal from the
experiment's path:

tensorboard --logdir='.'

(Alternatively, run the command from anywhere with the experiment's full path).

SuperGradients logs many useful metrics to tensorboard, including CPU and GPU usage, learning rate
scheduling, training and validation losses and other metrics, and many more. For example, let's check how the training
process goes by looking at the training's custom loss value:

<img src="images/external_model_example_train_loss.png" width="700">

We can also check the validation set's IoU metric's value:

<img src="images/external_model_example_valid_iou.png" width="700">

7. Predictions with the trained model Copy Chapter

Now that we have a trained model we can use it to make predictions on the test set. First, let's instantiate a test
dataset:

test_dataset = AISegmentDataset(data_path=data_path, split='test')

By instantiating the test dataset, all required pre-processing is already handled for us. Let's choose a single sample
from the dataset:

image, mask = test_dataset[0]
image, mask = image.unsqueeze(0), mask.unsqueeze(0)

Notice that we have added the batch dimension to the tensors. Next, we set the model to evaluation mode, and obtain
the predicted mask. Since our model learned to predict soft masks, we apply a threshold of 0.5 to binarize the mask:

model.eval()
pred = model(image)
pred = torch.gt(pred, 0.5).long()

Finally, let's apply the predicted mask to the image:

masked_image = image*pred

Now let's visualize the image, mask, and masked image to see how our model performed:

figure = plt.figure()
figure.add_subplot(1, 3, 1)
plt.title("Image")
plt.axis("off")
plt.imshow(image.detach().squeeze(0).permute(1, 2, 0))
figure.add_subplot(1, 3, 2)
plt.title("Mask")
plt.axis("off")
plt.imshow(pred.detach().squeeze(0).permute(1, 2, 0), cmap='gray')
figure.add_subplot(1, 3, 3)
plt.title("Masked Image")
plt.axis("off")
plt.imshow(masked_image.detach().squeeze(0).permute(1, 2, 0))
plt.show()

<img src="images/external_model_example_prediction.png" width="700">

---

Documentation/Source/Experiment Management (documentation/source/experiment_management.md) Copy Section

Experiment Management

Outline Copy Chapter

Checkpoint Root Directory (ckpt_root_dir) )')" title="Copy section prompt for LLMs"> Copy Section

Experiments (experiment_name) )')" title="Copy section prompt for LLMs"> Copy Section

Runs (run_id) )')" title="Copy section prompt for LLMs"> Copy Section

<ckpt_root_dir>
│
├── <experiment_name>
│   │
│   ├─── <run_dir>
│   │     ├─ ckpt_best.pth                   # Best performance during validation
│   │     ├─ ckpt_latest.pth                 # End of the most recent epoch
│   │     ├─ average_model.pth               # Averaged over specified epochs
│   │     ├─ ckpt_epoch_*.pth                # Checkpoints from certain epochs (e.g., epoch 10, 15)
│   │     ├─ events.out.tfevents.*           # Tensorflow run artifacts
│   │     └─ log_<timestamp>.txt             # Trainer logs of that particular run
│   │
│   └─── <other_run_dir>
│        └─ ...
│
└─── <other_experiment_name>
    │
    ├─── <run_dir>
    │     └─ ...
    │
    └─── <another_run_dir>
          └─ ...

#### A. Get the absolute path of a run directory
Manually navigate using <ckpt_root_dir>/<experiment_name>/<run_dir> or utilize the following programmatic approach:


from super_gradients.common.environment.checkpoints_dir_utils import get_checkpoints_dir_path
checkpoints_dir_path = get_checkpoints_dir_path(experiment_name="<experiment_name>", run_id="<run_id>")

#### B. Get the latest run id

from super_gradients.common.environment.checkpoints_dir_utils import get_latest_run_id
run_id = get_latest_run_id(experiment_name="<experiment_name>")

- Install dagshub and mlflow
- You can set up DagsHub according to the official documentation, or you'll be guided interactively to sign in when you run the code with the logger
- Adapt your code like in the following example

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,
      }
}
trainer.train(model=model, training_params=training_params, ...)

Weights & Biases Copy Section

- Install wandb
- Set up wandb according to the official documentation
- Make sure to login (You can check if you have a ~/.netrc token)
- Adapt your code like in the following example

from super_gradients import Trainer
trainer = Trainer("experiment_name")
model = ...
training_params = {
    ...                             # Your training params
    "sg_logger": "wandb_sg_logger", # Weights&Biases Logger, see class super_gradients.common.sg_loggers.wandb_sg_logger.WandBSGLogger for details
    "sg_logger_params":             # Params that will be passes to __init__ of the logger super_gradients.common.sg_loggers.wandb_sg_logger.WandBSGLogger
      {
        "project_name": "project_name", # W&B project name
        "save_checkpoints_remote": True,
        "save_tensorboard_remote": True,
        "save_logs_remote": True,
        "entity": "<YOUR-ENTITY-NAME>",         # username or team name where you're sending runs
        "api_server": "<OPTIONAL-WANDB-URL>"    # Optional: In case your experiment tracking is not hosted at wandb servers
      }
}
trainer.train(model=model, training_params=training_params, ...)

ClearML Copy Section

- Install clearml
- Set up CleaML according to the official documentation
- Adapt your code like in the following example

from super_gradients import Trainer
trainer = Trainer("experiment_name")
model = ...
training_params = {
    ...                                 # Your training params
    "sg_logger": "clearml_sg_logger",   # ClearML Logger, see class super_gradients.common.sg_loggers.wandb_sg_logger.ClearMLSGLogger for details
    "sg_logger_params":                 # Params that will be passes to __init__ of the logger super_gradients.common.sg_loggers.wandb_sg_logger.ClearMLSGLogger 
      {
        "project_name": "project_name", # ClearML project name
        "save_checkpoints_remote": True,
        "save_tensorboard_remote": True,
        "save_logs_remote": True,
      } 
}
trainer.train(model=model, training_params=training_params, ...)

Integrate any other Monitoring tool Copy Section

/ Detailed source-code truncated for AI context efficiency. /

from super_gradients import Trainer
trainer = Trainer("experiment_name")
model = ...
training_params = {
    ...,                                                    # Your training params
    "sg_logger": "CustomSGLogger",                          # Your custom CustomSGLogger
    "sg_logger_params": {"project_name": "my_project_name"} # Params that will be passed to __init__ of your CustomSGLogger  
}
trainer.train(model=model, training_params=training_params, ...)

- @multi_process_safe prevents multiple training nodes to do the same action. Check out DDP documentation for more details
- @register_logger() registers your class into our factory, allowing it to be instantiated from a string.
- sg_logger_params only requires project_name, the rest is provided by the Trainer.

Uploading custom objects with a callback Copy Chapter

from super_gradients.training.utils.callbacks.base_callbacks import PhaseContext, Callback
def do_something(inputs, target, preds):
    pass
class DetectionVisualizationCallback2(Callback):
    """Save a custom metric to tensorboard and wandb/clearml"""
    def __init__(self):
        super(Callback, self).__init__()
    def on_validation_batch_end(self, context: PhaseContext) -> None:
        # Do something using the PhaseContext
        custom_metric = do_something(context.inputs, context.target, context.preds)
        
        # Save it to the tensorboard and wandb/clearml
        context.sg_logger.add_scalar(
            tag="custom_metric",
            scalar_value=custom_metric,
            global_step=context.epoch,
        )

We encourage you to check out the API documentation of super_gradients.common.sg_loggers.base_sg_logger.BaseSGLogger to see every available method.

Chose your monitoring tool in the recipes Copy Chapter

sg_logger: wandb_sg_logger, # Weights&Biases Logger, see class super_gradients.common.sg_loggers.wandb_sg_logger.WandBSGLogger for details
sg_logger_params:             # Params that will be passes to __init__ of the logger super_gradients.common.sg_loggers.wandb_sg_logger.WandBSGLogger
  project_name: project_name, # W&B project name
  save_checkpoints_remote: True,
  save_tensorboard_remote: True,
  save_logs_remote: True,
  entity: <YOUR-ENTITY-NAME>,         # username or team name where you're sending runs
  api_server: <OPTIONAL-WANDB-URL>    # Optional: In case your experiment tracking is not hosted at wandb servers

To use AMP in SuperGradients, you simply need to set mixed_precision=True in your training_hyperparams.

In python script


from super_gradients import Trainer
trainer = Trainer("experiment_name")
model = ...
training_hyperparams = {"mixed_precision": True, ...:...}
trainer.train(model=model, training_hyperparams=training_hyperparams, ...)

In recipe


my_training_hyperparams.yaml
mixed_precision: True # Whether to use mixed precision or not.

Torch Compile Copy Chapter

PyTorch 2.0 introduced new torch.compile API which can be used to improve the training time of the model.
This API can be used to fuse the operations in the model graph and optimize the model for the target device.

SuperGradients support the torch.compile API and can be used to improve the training time of the model.
Here we report the relative improvement (reduction) of training time for several models and tasks. We measure the training time of one epoch.
This includes iteration over training and validation datasets, loss & metric computation. Essentially all steps that are performed during training.
Please note that the improvement vary depending on the model architecture, dataset, and training hyperparameters.

| Task | Recipe | Baseline (1 GPU) | Baseline (8 GPU) | 1 GPU With Compile | 8 GPU With Compile | Improvement, % (1 GPU) | Improvement, % (8 GPU) |
|-----------------------|---------------------------------|------------------|------------------|--------------------|--------------------|------------------------|------------------------|
| Semantic Segmentation | cityscapes_pplite_seg75 | 270.63 | 49.95 | 119.11 | 35.91 | 56% | 18% |
| Semantic Segmentation | cityscapes_regseg48 | 125.14 | 44.959 | 108.57 | 44.55 | 13.2% | 0.9% |
| Semantic Segmentation | cityscapes_segformer | 199.97 | 46.21 | 162.52 | 43.71 | 18.7% | 5.4% |
| Semantic Segmentation | cityscapes_stdc_seg75 | 425.19 | 73.07 | 153.16 | 45.89 | 63.9% | 37.19% |
| Semantic Segmentation | cityscapes_ddrnet | 226.51 | 51.78 | 174.11 | 48.29 | 23.1% | 7.3% |
| | | | | | | | |
| Object Detection | coco2017_yolo_nas_s | 1509 | 384.41 | 1379 | 376.10 | 8.6% | 2.42% |
| Object Detection | coco2017_yolo_nas_m | 2363 | 537.24 | 2090 | 508.40 | 11.5% | 0.19% |
| Object Detection | coco2017_yolo_nas_l | 3193 | 764.17 | 2869 | 745.58 | 10.14% | 2.43% |
| Object Detection | coco2017_ppyoloe_s/m/l/x | N/A | N/A | | N/A | N/A | N/A |
| Object Detection | coco2017_yolox_n/t/s/m/l/x | N/A | N/A | | N/A | N/A | N/A |
| Object Detection | coco2017_ssd_lite_mobilenet_v2 | N/A | N/A | | N/A | N/A | N/A |
| | | | | | | | |
| Classification | imagenet_efficientnet | | 425.61 | | 408.39 | | 4.1% |
| Classification | imagenet_mobilenetv3_large | | 373.73 | | 374.51 | | -0.2% |
| Classification | imagenet_regnetY | | 406.86 | | 383.04 | | 5.8% |
| Classification | imagenet_repvgg | | 407.19 | | 387.00 | | 4.9% |
| Classification | imagenet_resnet50 | | 481.36 | | 480.29 | | 0.22% |
| Classification | imagenet_vit_base | N/A | N/A | N/A | N/A | N/A | N/A |
| Classification | imagenet_vit_large | N/A | N/A | N/A | N/A | N/A | N/A |

In the table above, number are reported as speedup compared to the baseline training time.
Both experiments were run on 8x 3090 GPUs using PyTorch 2.0 with CUDA 11.8.
Training was done for 5 epochs and median value was picked to compute the speedup.
All experiments conducted with mixed precision (AMP) enabled, and SyncBN and EMA disabled.
Improvement percentage computed as follows: 100 * (baseline_time - compile_time) / baseline_time.

To leverage use of compiled models in SuperGradients one need to pass the torch_compile: True option to training hyperparameters:

python -m super_gradients.train_from_recipe --config-name=... training_hyperparams.torch_compile=True

In the YAML recipe:

my_training_recipe.yaml

training_hyperparams:
  torch_compile: True 
  torch_compile_mode: default | reduce-overhead | max-autotune


Or programmatically:

from super_gradients.training import Trainer
trainer = Trainer(
    ...,
    training_hyperparams = {
        "torch_compile": True,
        ...
        }
    )

Avoiding common pitfalls: Copy Section

* Don't use EMA with torch.compile.
* Don't use SyncBN with torch.compile.
* You may need to reduce batch size during training by quite a lot (Up to 2x)
* Training with mixed precision gives the best performance boost.

#### Exponential moving average EMA and Torch Compile

Torch Compile is still in its early stages and has some limitations. Not every model can be compiled.
Additionally, some training features can conflict with torch.compile.
Here is what we found so far:

* Exponential moving average EMA is incompatible with torch.compile (At the moment of writing, this is true for SG release 3.1.2).
If you want to use torch.compile in your training, you need to disable EMA when using torch.compile.


training_hyperparams:
    torch_compile: True 
    ema: False

When training using DDP, SyncBatchNorm layers _can be used_ with torch.compile simultaneously.
Unfortunately, due to implementation details the Sync BN have to break the model graph each time BN layer is encountered to perform BN sync operation.
That means you will most likely get no speedup from torch.compile when using Sync BN.
It also was observed that it may require more GPU memory compared to training without torch.compile.
If you are running into errors when using torch.compile and SyncBatchNorm simultaneously, you can try lowering the batch size to see if that helps.

#### Increased GPU memory consumption and Torch Compile

If during torch.compile you are getting wierd CUDA-related exception messages you can try reducing batch size.
When using reduce-overhead or max-autotune modes peak GPU memory consumption may be higher compared to training without torch.compile.
It is good idea to run code with CUDA_LAUNCH_BLOCKING=1 first as it usually provides more meaningful error messages.

#### Auto Mixed Precision and Torch Compile

Best speedup was achieved with combination of torch.compile and AMP enabled. For F32 training torch.compile may not provide any speedup at all.
This is highly dependent on the target GPU and support of fp32 tensor cores, so we leave it up to the user to decide whether to use AMP or not.

For AMP training with torch.compile enabled, you need to pass mixed_precision: True to training hyperparameters:


training_hyperparams:
    torch_compile: True 
    mixed_precision: True

---

Documentation/Source/Index (documentation/source/index.rst) Copy Section

.. SuperGradients documentation master file, created by
sphinx-quickstart on Wed Nov 24 10:50:24 2021.
You can adapt this file completely to your liking, but it should at least
contain the root toctree directive.

Welcome to SuperGradients's documentation!
==========================================