super-gradients

Why use DDP ?

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

How does it work ?

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

How to use it ?

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

main.py

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

# Initialize the environment
init_trainer()

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

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

# Everything you do below will run on 4 gpus

...

Trainer.train(...)

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

python main.py

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

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

torchrun --nproc_per_node=4 main.py

Calling functions on a single node

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

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

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

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

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

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

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

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

Good to know

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

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

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}

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": "/", # 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,
                     }
                   }

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

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

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)

• Python 3.7, 3.8 or 3.9 installed.
• 1.9.0 <= torch < 1.14
  • https://pytorch.org/get-started/locally/
• The python packages that are specified in requirements.txt;

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

See in PyPi

pip install super-gradients

That's it !

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