为什么使用 DDP？

近年来，深度学习模型变得越来越大，以至于在单个 GPU 上训练可能需要数周时间。 为了及时训练模型，必须使用多个 GPU 进行训练。 使用数百个 GPU 可以将模型的训练时间从一周缩短到不到一小时。

它是如何工作的？

每个 GPU 都有自己的进程，该进程控制着模型的一个副本，并从磁盘加载自己的小批量数据，在训练过程中将其发送到对应的 GPU。在每个 GPU 上完成前向传播后，梯度会在所有 GPU 之间进行归约，从而使所有 GPU 在本地拥有相同的梯度。这导致在反向传播之后，所有 GPU 上的模型权重保持同步。

如何使用它？

您只需几行代码就可以使用 SuperGradients 通过 DDP 训练您的模型。

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

# 初始化环境
init_trainer()

# 在 4 个 GPU 上启动 DDP
setup_device(multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL, num_gpus=4)

# 调用训练器
Trainer(expriment_name=...)

# 您在下方执行的所有操作都将在4个GPU上运行

...

Trainer.train(...)

最后，您可以通过一个简单的Python命令启动分布式训练。

python main.py

请注意，如果您使用的是torch<1.9.0（已弃用），则必须使用torch.distributed.launch或torchrun来启动训练，在这种情况下，nproc_per_node会覆盖通过gpu_mode设置的值：

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

torchrun --nproc_per_node=4 main.py

在单节点上调用函数

在DDP训练中，我们通常希望在主进程（即rank 0）上执行代码。在SG中，用户通常通过触发“阶段回调”来执行自己的代码（参见下文“使用阶段回调”部分）。可以使用ddp_silent_mode或multi_process_safe装饰器来确保所需代码仅在rank 0上运行。例如，考虑下面的简单阶段回调，它会在训练过程中将每个批次的前3张图像上传到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)

@multi_process_safe装饰器确保该回调仅由rank 0触发。或者，也可以通过SG训练器的布尔属性ddp_silent_mode来实现，该属性在当前进程排名为零时设置为False（即使在进程组被终止后仍然如此）：

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)

请注意，可以通过SgTrainer.ddp_silent_mode访问ddp_silent_mode。因此，在调用SgTrainer.train()之后，如果脚本中的某些部分只需要在rank 0上运行，则可以使用它。

需要了解的事项

您的总批量大小将是（GPU数量×批量大小），因此您可能需要增加学习率。虽然没有明确的规则，但一个经验法则是根据GPU数量线性增加学习率。

from super_gradients.training import models

# 实例化默认的预训练resnet18
default_resnet18 = models.get(model_name="resnet18", num_classes=100, pretrained_weights="imagenet")

# 实例化预训练resnet18，开启DropPath，概率为0.5
droppath_resnet18 = models.get(model_name="resnet18", arch_params={"droppath_prob": 0.5}, num_classes=100, pretrained_weights="imagenet")

# 实例化预训练resnet18，不带分类头。输出将来自全局池化之前的最后一个阶段
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

# 定义PyTorch的训练和验证数据加载器以及优化器

# 定义回调中要调用的内容
rop_lr_scheduler = ReduceLROnPlateau(optimizer, mode="max", patience=10, verbose=True)

# 定义阶段回调，它们将按照Phase中定义的方式触发
phase_callbacks = [LRSchedulerCallback(scheduler=rop_lr_scheduler,
                                       phase=Phase.VALIDATION_EPOCH_END,
                                       metric_name="Accuracy")]

# 创建一个训练器对象，更多参数请参阅声明
trainer = Trainer("experiment_name")

# 将phase_callbacks定义为训练参数的一部分
train_params = {"phase_callbacks": phase_callbacks}

from super_gradients import Trainer

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

training_params = { ...  # 您的训练参数
                   "sg_logger": "dagshub_sg_logger",  # DagsHub日志记录器，详情请参阅super_gradients.common.sg_loggers.dagshub_sg_logger.DagsHubSGLogger类
                   "sg_logger_params":  # 将传递给日志记录器super_gradients.common.sg_loggers.dagshub_sg_logger.DagsHubSGLogger初始化方法的参数
                     {
                       "dagshub_repository": "<REPO_OWNER>/<REPO_NAME>", # 可选：您的DagsHub项目名称，由所有者名称、斜杠和仓库名称组成。如果留空，您将在运行时被提示手动填写。
                       "log_mlflow_only": False, # 可选：设置为true以跳过DVC日志记录，仅将所有工件记录到MLflow  
                       "save_checkpoints_remote": True,
                       "save_tensorboard_remote": True,
                       "save_logs_remote": True,
                     }
                   }

from super_gradients import Trainer

# 创建一个训练器对象，更多参数请查看声明
trainer = Trainer("experiment_name")

train_params = { ... # 训练参数
                "sg_logger": "wandb_sg_logger", # Weights&Biases 日志记录器，详情请参阅 WandBSGLogger 类
                "sg_logger_params": # 将传递给日志记录器 __init__ 方法的参数 
                  {
                    "project_name": "project_name", # W&B 项目名称
                    "save_checkpoints_remote": True,
                    "save_tensorboard_remote": True,
                    "save_logs_remote": True,
                  } 
               }

from super_gradients import Trainer

# 创建一个训练器对象，更多参数请查看声明
trainer = Trainer("experiment_name")

train_params = { ... # 训练参数
                "sg_logger": "clearml_sg_logger", # ClearML 日志记录器，详情请参阅 ClearMLSGLogger 类
                "sg_logger_params": # 将传递给日志记录器 __init__ 方法的参数 
                  {
                    "project_name": "project_name", # ClearML 项目名称
                    "save_checkpoints_remote": True,
                    "save_tensorboard_remote": True,
                    "save_logs_remote": True,
                  } 
               }

您可以使用 apply_model() 方法将 SuperGradients YOLO-NAS 模型直接应用于您的 FiftyOne 数据集：

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)

SuperGradients YOLO-NAS 模型可以直接从 FiftyOne Model Zoo 获取：

import fiftyone as fo
import五十one.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 或 3.9。
• 1.9.0 <= torch < 1.14
  • https://pytorch.org/get-started/locally/
• requirements.txt 中指定的 Python 包；

• NVIDIA CUDA Toolkit >= 11.2
• CuDNN >= 8.1.x
• 支持 CUDA >= 11.2 的 NVIDIA 驱动程序（≥460.x）

请参阅 PyPI

pip install super-gradients

就完成了！

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