PyTorch多卡分布式训练DistributedDataParallel 使用方法

2023-05-09 17:21:04 买帖  | 投诉/举报

篇首语:本文由小编为大家整理,主要介绍了PyTorch多卡分布式训练DistributedDataParallel 使用方法相关的知识,希望对你有一定的参考价值。

PyTorch多卡分布式训练DistributedDataParallel 使用方法

目录

PyTorch多卡分布式训练DistributedDataParallel 使用方法

1.DP模式和DP模式

(1)单进程多GPU训练模式:DP模式

(2)多进程多GPU训练模式:DDP模式

2.Pytorch分布式训练方法

3.Pytorch-Base-Trainer(PBT)分布式训练工具

(1) 工具介绍

(2) 安装

(3)使用方法

4.Example: 构建自己的分类Pipeline

5.可视化


前言:

  •  开源不易,麻烦给个【Star】
  • Github: GitHub - PanJinquan/Pytorch-Base-Trainer: Pytorch分布式训练框架
  • pip安装包: basetrainer · PyPI
  • 博客地址:Pytorch基础训练库Pytorch-Base-Trainer(支持模型剪枝 分布式训练)_pan_jinquan的博客-CSDN博客

【尊重原则,转载请注明出处】:https://panjinquan.blog.csdn.net/article/details/122702287

基于Pytorch-Base-Trainer(PBT)分布式训练工具,已经形成了一套完整的android端上部署流程,支持CPU和GPU

人体姿态估计2DPose人脸+人体检测人像抠图
CPU/GPU:70/50msCPU/GPU:30/20msCPU/GPU:150/30ms

1.DP模式和DP模式

Pytorch多卡训练有两种方式,一种是单进程多GPU训练模式(single process multi-gpus),另一种的多进程多卡模式(multi-processes multi-gpus)

(1)单进程多GPU训练模式:DP模式

Pytorch通过nn.DataParallel可实现多卡训练模型(简称DP模式),这是single process multi-gpus 的多卡并行机制,这种并行模式下并行的多卡都是由一个进程进行控制,其缺点有:

  1. 尽管 DataLoader 可以指定 num_worker,增加负责加载数据的线程数量,但是线程的资源受限于父进程,且由于python的GIL机制,不能利用好多核的并行优势
  2. 模型在 gpu 群组中进行初始化与广播过程依赖单一进程的串行操作
  3. DP模式相当于将多个GPU卡合并为一个卡进行训练

尽管DataParallel更易于使用(只需简单包装单GPU模型),但由于使用一个进程来计算模型权重,然后在每个批处理期间将分发到每个GPU,因此通信很快成为一个瓶颈,GPU利用率通常很低。而且,nn.DataParallel要求所有的GPU都在同一个节点上(不支持分布式)。

(2)多进程多GPU训练模式:DDP模式

Pytorch通过nn.parallel.DistributedDataParallel可实现多进程多卡训练模型(也称DDP模式),这种多卡并行机制的特点/优势有:

  1. 一个进程一个GPU(当然可以让每个进程控制多个GPU,但这显然比每个进程有一个GPU要慢)
  2. 充分利用多核并行的优势加载数据
  3. 模型在 gpu 群组中进行初始化的过程由各自的进程负责调度
  4. 代码可以无缝切换单机多卡与多机多卡训练,因为此时单机单卡成为了单机多卡/多机多卡并行下的一个特例
  5. GPU可以都在同一个节点上,也可以分布在多个节点上。每个进程都执行相同的任务,并且每个进程都与所有其他进程通信。进程或者说GPU之间只传递梯度,这样网络通信就不再是瓶颈。

在训练过程中,每个进程从磁盘加载batch数据,并将它们传递到其GPU。每一个GPU都有自己的前向过程,然后梯度在各个GPUs间进行All-Reduce。每一层的梯度不依赖于前一层,所以梯度的All-Reduce和后向过程同时计算,以进一步缓解网络瓶颈。在后向过程的最后,每个节点都得到了平均梯度,这样模型参数保持同步。

这就要求多个进程,甚至多个节点上的多个进程实现同步并通信。Pytorch通过distributed.init_process_group函数来实现这一点。他需要知道进程0位置以便所有进程都可以同步,以及预期的进程总数。每个进程都需要知道进程总数及其在进程中的顺序,以及使用哪个GPU。通常将进程总数称为world_size。 

Pytorch提供了nn.utils.data.DistributedSampler来为各个进程切分数据,以保证训练数据不重叠。

nn.DataParallelnn.distributedataparallel的主要差异可以总结为以下几点:

  1. DistributedDataParallel支持模型并行,而DataParallel并不支持,这意味如果模型太大单卡显存不足时只能使用前者;
  2. DataParallel是单进程多线程的,只用于单机情况,而DistributedDataParallel是多进程的,适用于单机和多机情况,真正实现分布式训练;
  3. DistributedDataParallel的训练更高效,因为每个进程都是独立的Python解释器,避免GIL问题,而且通信成本低其训练速度更快,基本上DataParallel已经被弃用;
  4. 必须要说明的是DistributedDataParallel中每个进程都有独立的优化器,执行自己的更新过程,但是梯度通过通信传递到每个进程,所有执行的内容是相同的;

除了PyTorch官方实现的分布式训练方案,还有horovod分布式训练工具,不仅支持PyTorch还支持TensorFlow和MXNet框架,实现起来也是比较容易的,速度方面应该不相上下。

参考资料:PyTorch分布式训练简明教程 - 知乎


2.Pytorch分布式训练方法

分布式训练一般分为数据并行模型并行两种,Pytorch分布式训练的实现步骤可简述如下:

  1. 首先在nn.DataParallel(即DP模式下)实现多卡加载数据,训练模型并调试成功;这一步是为了保证你的训练流程正常无BUG。然后就可以开始魔改了
  2.  数据并行(分布式)DataLoader的样本采样器(sampler)修改为分布式采样器torch_utils.distributed.DistributedSampler
  3. 模型并行(分布式):将torch.nn.parallel.DistributedDataParallel 代替torch.nn.DataParallel 
  4. 为了能够使用 DistributedDataParallel 需要先进行进程间通讯环境的初始化,torch.distributed.init_process_group()
  5. 为了解决并行训练中加载到各个 worker/gpu 中的 sub-mini-batch 之间出现 example overlap 问题,还可以配合 torch.utils.data.distributed.DistributedSampler 进行使用
  6. 为了让进程与 gpu 进行一一匹配,在程序的开头通过 torch.cuda.set_device 设定目标设备
  7. (可选)为了让各个 worker/gpu 能有一致的初始值,在程序开头通过 torch.manual_seed 与 torch.cuda.manual_seed 来初始化随机数种子

所以代码结构如下:

# filename: distributed_example.py# import some module...... parser = argparse.Argument()parser.add_argument('--init_method', defalut='env://', type=str)parser.add_argument('--local_rank', type=int, default=0)args = parser.parse() import os# Set master information and NIC# NIC for communicationos.environ['NCCL_SOCKET_IFNAME'] = 'xxxx'# set master node address# recommend setting to ib NIC to gain bigger communication bandwidthos.environ['MASTER_ADDR'] = '192.168.xx.xx'# set master node port# **caution**: avoid port conflictos.environ['MASTER_PORT'] = '1234' def main():    # step 1    # set random seed    torch.manual_seed(seed)    torch.cuda.manual_seed(seed)         # step 2    # set target device    torch.cuda.set_device(args.local_rank)      # step 3    # initialize process group    # use nccl backend to speedup gpu communication    torch.distributed.init_process_group(backend='nccl', init_method=args.init_method)         ...    ...    # step 4    # set distributed sampler    # the same, you can set distributed sampler for validation set    train_sampler = torch.utils.data.distributed.DistributedSampler(        dataset_train)    train_loader = torch.utils.data.DataLoader(        dataset_train, batch_size=BATCH_SIZE, sampler=train_sampler, pin_memory=PIN_MEMORY,        num_workers=NUM_WORKERS    )         ...    ...    # step 5    # initialize model     model = resnet50()    model.cuda()         # step 6    # wrap model with distributeddataparallel    # map device with model process, and we bind process n with gpu n(n=0,1,2,3...) by default.    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)         ...    ...    for epoch in range(epochs):        # step 7        # update sampler with epoch        train_sampler.set_epoch(epoch)                 # then do whatever you want to do, just like a single device/gpu training program        ...        ...

注意:

       由于程序开头采用了 torch.cuda.set_device() 指定了目标的设备,所以后续的操作中如果有涉及要将数据、模型移动到 gpu 中的操作需要改为 model.cuda()、inputs.cuda(),该指令会将对象正确地复制到对应的 cuda 设备中。

如果你执意使用 to 操作,那么注意确保 xx.to('cuda:n') 中的 cuda:n 与目标设备是匹配的。

以多进程方式启动训练脚本

你当然可以以 python distributed_example.py 这样的形式启动训练脚本,不过这样无法触发多进程。pytorch 为多进程启动脚本提供了 launch 工具,所以正确的启动方式是:

python -m torch.distributed.launch --nnodes=<nodes> --nproc_per_node=<process per node> --node_rank=<rank of current node>\\        distributed_example.py --arg1 --arg2 and all other arguments of your trainning script

 参数说明:

  • nnodes:指定参与计算的节点数量,默认值为1,单机多卡的训练中可以不用指定
  • nproc_per_node:指定每个节点中的所要启动的进程数量,由于进程与 gpu 一一对应,所以这里的数值不能大于系统中所能使用的 gpu 数量
  • node_rank: 指定当前节点在整个系统中的序号,从 0 开始递增,需要注意的是,在多机多卡训练中 node_rank == 0 的节点表示 master,所以 node_rank == 0 的节点必须是 MASTER_ADDR 所在的节点,否则多卡间的通信无法正确建立连接。

老实说,【从DP模式升级到DDP的方法】看起来简单,步骤也不多,但真正要跑起来还是很多地方需要优化的;

这种多进程训练的方法,每个进程需要分配一个卡进行训练,这就导致你保存模型,打印Log,测试数据都变成复杂了,比如会出现多个进程都会打印相同的Log的问题;一般建议你,定义一个主进程,且在主进程中打印Log,保存模型,测试数据等操作,这样可以避免上述问题了

那有没有一个简单方法,可以快速实现Pytorch的分布式训练

有的,我今天就介绍一个我自己整合的Pytorch的分布式训练工具Pytorch-Base-Trainer,基于这套工具,你可以简单配置,即可实现DP或者DDP模式的训练,而无需关注各种进程间通讯,端口设置等这些复杂的过程。


3.Pytorch-Base-Trainer(PBT)分布式训练工具

(1) 工具介绍

考虑到深度学习训练过程都有一套约定成俗的流程,鄙人借鉴Keras开发了一套基础训练库: Pytorch-Base-Trainer(PBT); 这是一个基于Pytorch开发的基础训练库,支持以下特征:

  •  支持多卡训练训练(DP模式)和分布式多卡训练(DDP模式),参考build_model_parallel
  •  支持argparse命令行指定参数,也支持config.yaml配置文件
  •  支持最优模型保存ModelCheckpoint
  •  支持自定义回调函数Callback
  •  支持NNI模型剪枝(L1/L2-Pruner,FPGM-Pruner Slim-Pruner)nni_pruning
  •  非常轻便,安装简单

博客介绍:

Pytorch基础训练库Pytorch-Base-Trainer(支持模型剪枝 分布式训练)_pan_jinquan的博客-CSDN博客考虑到深度学习训练过程都有一套约定成俗的流程,鄙人借鉴Keras开发了一套基础训练库: Pytorch-Base-Trainer(PBT); 这是一个基于Pytorch开发的基础训练库,支持以下特征:https://panjinquan.blog.csdn.net/article/details/122662902GitHub地址:

GitHub - PanJinquan/Pytorch-Base-Trainer: Pytorch分布式训练框架https://github.com/PanJinquan/Pytorch-Base-Trainer

(2) 安装

  • 源码安装
git clone https://github.com/PanJinquan/Pytorch-Base-Trainercd Pytorch-Base-Trainerbash setup.sh #pip install dist/basetrainer-*.*.*.tar.gz
  • pip安装
pip install basetrainer
  • 使用NNI 模型剪枝工具,需要安装NNI
# Linux or macOSpython3 -m pip install --upgrade nni# Windowspython -m pip install --upgrade nni

(3)使用方法

basetrainer使用方法可以参考example.py,构建自己的训练器,可通过如下步骤实现:

  • step1: 新建一个类ClassificationTrainer,继承trainer.EngineTrainer
  • step2: 实现接口
def build_train_loader(self, cfg, **kwargs):    """定义训练数据"""    raise NotImplementedError("build_train_loader not implemented!")def build_test_loader(self, cfg, **kwargs):    """定义测试数据"""    raise NotImplementedError("build_test_loader not implemented!")def build_model(self, cfg, **kwargs):    """定于训练模型"""    raise NotImplementedError("build_model not implemented!")def build_optimizer(self, cfg, **kwargs):    """定义优化器"""    raise NotImplementedError("build_optimizer not implemented!")def build_criterion(self, cfg, **kwargs):    """定义损失函数"""    raise NotImplementedError("build_criterion not implemented!")def build_callbacks(self, cfg, **kwargs):    """定义回调函数"""    raise NotImplementedError("build_callbacks not implemented!")step3: 在初始化中调用builddef __init__(self, cfg):    super(ClassificationTrainer, self).__init__(cfg)    ...    self.build(cfg)    ...step4: 实例化ClassificationTrainer,并使用launch启动分布式训练def main(cfg):    t = ClassificationTrainer(cfg)    return t.run()if __name__ == "__main__":    parser = get_parser()    args = parser.parse_args()    cfg = setup_config.parser_config(args)    launch(main,           num_gpus_per_machine=len(cfg.gpu_id),           dist_url="tcp://127.0.0.1:28661",           num_machines=1,           machine_rank=0,           distributed=cfg.distributed,           args=(cfg,))

4.Example: 构建自己的分类Pipeline

  • basetrainer使用方法可以参考example.py
# 单进程多卡训练python example.py --gpu_id 0 1 # 使用命令行参数python example.py --config_file configs/config.yaml # 使用yaml配置文件# 多进程多卡训练(分布式训练)python example.py --config_file configs/config.yaml --distributed # 使用yaml配置文件
  • 目标支持的backbone有:resnet[18,34,50,101], ,mobilenet_v2等,详见backbone等 ,其他backbone可以自定义添加
  • 训练参数可以通过两种方法指定: (1) 通过argparse命令行指定 (2)通过https://panjinquan.blog.csdn.net/article/details/122702287配置文件,当存在同名参数时,以配置文件为默认值
参数类型参考值说明
train_datastr, list-训练数据文件,可支持多个文件
test_datastr, list-测试数据文件,可支持多个文件
work_dirstrwork_space训练输出工作空间
net_typestrresnet18backbone类型,resnet,resnest,mobilenet_v2,...
input_sizelist[128,128]模型输入大小[W,H]
batch_sizeint32batch size
lrfloat0.1初始学习率大小
optim_typestrSGD优化器,SGD,Adam
loss_typestrCELoss损失函数
schedulerstrmulti-step学习率调整策略,multi-step,cosine
milestoneslist[30,80,100]降低学习率的节点,仅仅scheduler=multi-step有效
momentumfloat0.9SGD动量因子
num_epochsint120循环训练的次数
num_warn_upint3warn_up的次数
num_workersint12DataLoader开启线程数
weight_decayfloat5e-4权重衰减系数
gpu_idlist[ 0 ]指定训练的GPU卡号,可指定多个
log_freqin20显示LOG信息的频率
finetunestrmodel.pthfinetune的模型
use_pruneboolTrue是否进行模型剪枝
progressboolTrue是否显示进度条
distributedboolFalse是否使用分布式训练

 一个简单分类例子如下:

# -*-coding: utf-8 -*-"""    @Author : panjq    @E-mail : pan_jinquan@163.com    @Date   : 2021-07-28 22:09:32"""import osimport syssys.path.append(os.getcwd())import argparseimport basetrainerfrom torchvision import transformsfrom torchvision.datasets import ImageFolderfrom basetrainer.engine import trainerfrom basetrainer.engine.launch import launchfrom basetrainer.criterion.criterion import get_criterionfrom basetrainer.metric import accuracy_recorderfrom basetrainer.callbacks import log_history, model_checkpoint, losses_recorder, multi_losses_recorderfrom basetrainer.scheduler import build_schedulerfrom basetrainer.optimizer.build_optimizer import get_optimizerfrom basetrainer.utils import log, file_utils, setup_config, torch_toolsfrom basetrainer.models import build_modelsprint(basetrainer.__version__)class ClassificationTrainer(trainer.EngineTrainer):    """ Training Pipeline """    def __init__(self, cfg):        super(ClassificationTrainer, self).__init__(cfg)        torch_tools.set_env_random_seed()        cfg.model_root = os.path.join(cfg.work_dir, "model")        cfg.log_root = os.path.join(cfg.work_dir, "log")        if self.is_main_process:            file_utils.create_dir(cfg.work_dir)            file_utils.create_dir(cfg.model_root)            file_utils.create_dir(cfg.log_root)            file_utils.copy_file_to_dir(cfg.config_file, cfg.work_dir)            setup_config.save_config(cfg, os.path.join(cfg.work_dir, "setup_config.yaml"))        self.logger = log.set_logger(level="debug",                                     logfile=os.path.join(cfg.log_root, "train.log"),                                     is_main_process=self.is_main_process)        # build project        self.build(cfg)        self.logger.info("=" * 60)        self.logger.info("work_dir          :".format(cfg.work_dir))        self.logger.info("config_file       :".format(cfg.config_file))        self.logger.info("gpu_id            :".format(cfg.gpu_id))        self.logger.info("main device       :".format(self.device))        self.logger.info("num_samples(train):".format(self.num_samples))        self.logger.info("num_classes       :".format(cfg.num_classes))        self.logger.info("mean_num          :".format(self.num_samples / cfg.num_classes))        self.logger.info("=" * 60)    def build_optimizer(self, cfg, **kwargs):        """build_optimizer"""        self.logger.info("build_optimizer")        self.logger.info("optim_type:,init_lr:,weight_decay:".format(cfg.optim_type, cfg.lr, cfg.weight_decay))        optimizer = get_optimizer(self.model,                                  optim_type=cfg.optim_type,                                  lr=cfg.lr,                                  momentum=cfg.momentum,                                  weight_decay=cfg.weight_decay)        return optimizer    def build_criterion(self, cfg, **kwargs):        """build_criterion"""        self.logger.info("build_criterion,loss_type:,num_classes:".format(cfg.loss_type, cfg.num_classes))        criterion = get_criterion(cfg.loss_type, cfg.num_classes, device=self.device)        return criterion    def build_train_loader(self, cfg, **kwargs):        """build_train_loader"""        self.logger.info("build_train_loader,input_size:".format(cfg.input_size))        transform = transforms.Compose([            transforms.Resize([int(128 * cfg.input_size[1] / 112), int(128 * cfg.input_size[0] / 112)]),            transforms.RandomHorizontalFlip(),            transforms.RandomCrop([cfg.input_size[1], cfg.input_size[0]]),            transforms.ToTensor(),            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),        ])        dataset = ImageFolder(root=cfg.train_data, transform=transform)        cfg.num_classes = len(dataset.classes)        cfg.classes = dataset.classes        loader = self.build_dataloader(dataset, cfg.batch_size, cfg.num_workers, phase="train",                                       shuffle=True, pin_memory=False, drop_last=True, distributed=cfg.distributed)        return loader    def build_test_loader(self, cfg, **kwargs):        """build_test_loader"""        self.logger.info("build_test_loader,input_size:".format(cfg.input_size))        transform = transforms.Compose([            transforms.Resize([int(128 * cfg.input_size[1] / 112), int(128 * cfg.input_size[0] / 112)]),            transforms.CenterCrop([cfg.input_size[1], cfg.input_size[0]]),            transforms.ToTensor(),            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),        ])        dataset = ImageFolder(root=cfg.train_data, transform=transform)        loader = self.build_dataloader(dataset, cfg.batch_size, cfg.num_workers, phase="test",                                       shuffle=False, pin_memory=False, drop_last=False, distributed=False)        return loader    def build_model(self, cfg, **kwargs):        """build_model"""        self.logger.info("build_model,net_type:".format(cfg.net_type))        model = build_models.get_models(net_type=cfg.net_type, input_size=cfg.input_size,                                        num_classes=cfg.num_classes, pretrained=True)        if cfg.finetune:            self.logger.info("finetune:".format(cfg.finetune))            state_dict = torch_tools.load_state_dict(cfg.finetune)            model.load_state_dict(state_dict)        if cfg.use_prune:            from basetrainer.pruning import nni_pruning            sparsity = 0.2            self.logger.info("use_prune:,sparsity:".format(cfg.use_prune, sparsity))            model = nni_pruning.model_pruning(model,                                              input_size=[1, 3, cfg.input_size[1], cfg.input_size[0]],                                              sparsity=sparsity,                                              reuse=False,                                              output_prune=os.path.join(cfg.work_dir, "prune"))        model = self.build_model_parallel(model, cfg.gpu_id, distributed=cfg.distributed)        return model    def build_callbacks(self, cfg, **kwargs):        """定义回调函数"""        self.logger.info("build_callbacks")        # 准确率记录回调函数        acc_record = accuracy_recorder.AccuracyRecorder(target_names=cfg.classes,                                                        indicator="Accuracy")        # loss记录回调函数        loss_record = losses_recorder.LossesRecorder(indicator="loss")        # Tensorboard Log等历史记录回调函数        history = log_history.LogHistory(log_dir=cfg.log_root,                                         log_freq=cfg.log_freq,                                         logger=self.logger,                                         indicators=["loss", "Accuracy"],                                         is_main_process=self.is_main_process)        # 模型保存回调函数        checkpointer = model_checkpoint.ModelCheckpoint(model=self.model,                                                        optimizer=self.optimizer,                                                        moder_dir=cfg.model_root,                                                        epochs=cfg.num_epochs,                                                        start_save=-1,                                                        indicator="Accuracy",                                                        logger=self.logger)        # 学习率调整策略回调函数        lr_scheduler = build_scheduler.get_scheduler(cfg.scheduler,                                                     optimizer=self.optimizer,                                                     lr_init=cfg.lr,                                                     num_epochs=cfg.num_epochs,                                                     num_steps=self.num_steps,                                                     milestones=cfg.milestones,                                                     num_warn_up=cfg.num_warn_up)        callbacks = [acc_record,                     loss_record,                     lr_scheduler,                     history,                     checkpointer]        return callbacks    def run(self, logs: dict = ):        self.logger.info("start train")        super().run(logs)def main(cfg):    t = ClassificationTrainer(cfg)    return t.run()def get_parser():    parser = argparse.ArgumentParser(description="Training Pipeline")    parser.add_argument("-c", "--config_file", help="configs file", default="configs/config.yaml", type=str)    # parser.add_argument("-c", "--config_file", help="configs file", default=None, type=str)    parser.add_argument("--train_data", help="train data", default="./data/dataset/train", type=str)    parser.add_argument("--test_data", help="test data", default="./data/dataset/val", type=str)    parser.add_argument("--work_dir", help="work_dir", default="output", type=str)    parser.add_argument("--input_size", help="input size", nargs="+", default=[224, 224], type=int)    parser.add_argument("--batch_size", help="batch_size", default=32, type=int)    parser.add_argument("--gpu_id", help="specify your GPU ids", nargs="+", default=[0], type=int)    parser.add_argument("--num_workers", help="num_workers", default=0, type=int)    parser.add_argument("--num_epochs", help="total epoch number", default=50, type=int)    parser.add_argument("--scheduler", help=" learning scheduler: multi-step,cosine", default="multi-step", type=str)    parser.add_argument("--milestones", help="epoch stages to decay learning rate", nargs="+",                        default=[10, 20, 40], type=int)    parser.add_argument("--num_warn_up", help="num_warn_up", default=3, type=int)    parser.add_argument("--net_type", help="net_type", default="mobilenet_v2", type=str)    parser.add_argument("--finetune", help="finetune model file", default=None, type=str)    parser.add_argument("--loss_type", help="loss_type", default="CELoss", type=str)    parser.add_argument("--optim_type", help="optim_type", default="SGD", type=str)    parser.add_argument("--lr", help="learning rate", default=0.1, type=float)    parser.add_argument("--weight_decay", help="weight_decay", default=0.0005, type=float)    parser.add_argument("--momentum", help="momentum", default=0.9, type=float)    parser.add_argument("--log_freq", help="log_freq", default=10, type=int)    parser.add_argument('--use_prune', action='store_true', help='use prune', default=False)    parser.add_argument('--progress', action='store_true', help='display progress bar', default=True)    parser.add_argument('--distributed', action='store_true', help='use distributed training', default=False)    parser.add_argument('--polyaxon', action='store_true', help='polyaxon', default=False)    return parserif __name__ == "__main__":    parser = get_parser()    cfg = setup_config.parser_config(parser.parse_args(), cfg_updata=True)    launch(main,           num_gpus_per_machine=len(cfg.gpu_id),           dist_url="tcp://127.0.0.1:28661",           num_machines=1,           machine_rank=0,           distributed=cfg.distributed,           args=(cfg,))

5.可视化

目前训练过程可视化工具是使用Tensorboard,使用方法:

tensorboard --logdir=path/to/log/

以上是关于PyTorch多卡分布式训练DistributedDataParallel 使用方法的主要内容,如果未能解决你的问题,请参考以下文章