Super Gradients用DDP训练Yolo NAS遇DistributedSampler属性错误

错误原因 #

从报错堆栈可以看出，在_instantiate_sampler函数中，代码试图调用dataloader_params["sampler"]的keys()方法，但此时sampler已经是实例化的DistributedSampler对象，而非框架预期的字典结构。这是因为调用setup_device(multi_gpu='DDP')后，PyTorch Elastic会重启脚本并启动多个子进程，每个子进程都会重新执行代码，导致数据加载器的sampler被框架自动注入后，又被重复处理，引发类型不匹配。

解决方案 #

将数据加载器的创建逻辑放到if __name__ == "__main__":代码块中，确保只有主进程执行数据加载器的初始化，避免子进程重复执行导致sampler冲突。同时移除手动将模型转移到DEVICE的代码，Trainer会自动处理设备分配。

修改后的完整代码 #

import torch
import os
from PIL import Image

from super_gradients.training import Trainer, dataloaders, models
from super_gradients.training.dataloaders.dataloaders import (
    coco_detection_yolo_format_train, coco_detection_yolo_format_val
)
from super_gradients.training.losses import PPYoloELoss
from super_gradients.training.metrics import DetectionMetrics_050
from super_gradients.training.models.detection_models.pp_yolo_e import (
    PPYoloEPostPredictionCallback
)
from super_gradients.training.utils.distributed_training_utils import setup_device

class config:
    # trainer params
    CHECKPOINT_DIR = 'checkpoints' # 指定 checkpoint 保存路径
    EXPERIMENT_NAME = 'experiment_v2' # 指定实验名称

    # dataset params
    DATA_DIR = 'yolo_data' # 数据集根目录

    TRAIN_IMAGES_DIR = 'train/images' # 训练集图片子目录
    TRAIN_LABELS_DIR = 'train/labels' # 训练集标签子目录

    VAL_IMAGES_DIR = 'val/images' # 验证集图片子目录
    VAL_LABELS_DIR = 'val/labels' # 验证集标签子目录

    # 测试集（如果有）
    TEST_IMAGES_DIR = 'test/images' # 测试集图片子目录
    TEST_LABELS_DIR = 'test/labels' # 测试集标签子目录

    CLASSES = ['Face'] # 类别列表
    NUM_CLASSES = len(CLASSES)

    # dataloader 参数
    DATALOADER_PARAMS={
        'batch_size':64,
        'num_workers':4
    }

    # model params
    MODEL_NAME = 'yolo_nas_l' # 可选 yolo_nas_s, yolo_nas_m, yolo_nas_l
    PRETRAINED_WEIGHTS = 'coco' # 预训练权重仅支持 coco
    DEVICE = 'cuda' if torch.cuda.is_available() else "cpu"

if __name__ == "__main__":
    setup_device(multi_gpu='DDP', num_gpus=2)
    trainer = Trainer(experiment_name=config.EXPERIMENT_NAME, ckpt_root_dir=config.CHECKPOINT_DIR)

    train_data = coco_detection_yolo_format_train(
        dataset_params={
            'data_dir': config.DATA_DIR,
            'images_dir': config.TRAIN_IMAGES_DIR,
            'labels_dir': config.TRAIN_LABELS_DIR,
            'classes': config.CLASSES
        },
        dataloader_params=config.DATALOADER_PARAMS
    )

    val_data = coco_detection_yolo_format_val(
        dataset_params={
            'data_dir': config.DATA_DIR,
            'images_dir': config.VAL_IMAGES_DIR,
            'labels_dir': config.VAL_LABELS_DIR,
            'classes': config.CLASSES
        },
        dataloader_params=config.DATALOADER_PARAMS
    )

    model = models.get(config.MODEL_NAME, 
                       num_classes=config.NUM_CLASSES, 
                       pretrained_weights=config.PRETRAINED_WEIGHTS
                       )

    train_params = {
        "average_best_models":True,
        "warmup_mode": "linear_epoch_step",
        "warmup_initial_lr": 8e-6,
        "lr_warmup_epochs": 5,
        "initial_lr": 40e-4,
        "lr_mode": "cosine",
        "cosine_final_lr_ratio": 0.1,
        "optimizer": "Adam",
        "optimizer_params": {"weight_decay": 0.0001},
        "zero_weight_decay_on_bias_and_bn": True,
        "ema": True,
        "ema_params": {"decay": 0.9, "decay_type": "threshold"},
        "max_epochs": 300,
        "mixed_precision": True,
        "loss": PPYoloELoss(
            use_static_assigner=False,
            num_classes=config.NUM_CLASSES,
            reg_max=16
        ),
        "valid_metrics_list": [
            DetectionMetrics_050(
                score_thres=0.1,
                top_k_predictions=300,
                num_cls=config.NUM_CLASSES,
                normalize_targets=True,
                post_prediction_callback=PPYoloEPostPredictionCallback(
                    score_threshold=0.01,
                    nms_top_k=1000,
                    max_predictions=300,
                    nms_threshold=0.7
                )
            )
        ],
        "metric_to_watch": 'mAP@0.50'
    }

    trainer.train(model=model, 
                  training_params=train_params, 
                  train_loader=train_data, 
                  valid_loader=val_data)

补充说明 #

• setup_device启动DDP时会通过PyTorch Elastic重启脚本，子进程会重新执行整个代码文件，if __name__ == "__main__"能确保数据加载器只在主进程初始化，避免子进程重复处理sampler。
• 移除model.to(config.DEVICE)是因为Super Gradients的Trainer会自动处理模型的设备分配，包括分布式场景下的多GPU分布。

内容的提问来源于stack exchange，提问作者Alim Tleuliyev