super-gradients

为什么使用DDP？

最近的深度学习模型越来越大，以至于在单个GPU上训练可能需要数周时间。 为了及时训练模型，有必要使用多个GPU进行训练。 使用100多个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 trainer 的布尔属性 ddp_silent_mode 来实现(phase context 可以访问该属性)。该属性仅当当前进程 rank 为零时设为 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)

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

值得了解

你的总批量大小将是 (GPU 数量 x 批量大小),所以你可能需要增加学习率。 没有明确的规则,但经验法则似乎是根据 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 的 __init__ 的参数
                     {
                       "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)

可以直接从 FiftyOne Model Zoo 访问 SuperGradients YOLO-NAS 模型:

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