PyTorch 分布式训练 (DP/DDP/torchrun/多机多卡) ＜笔记总结＞

xiangyong58

已于 2022-07-06 20:16:28 修改

阅读量4k

点赞数 2

CC 4.0 BY-SA版权

分类专栏： Machine & Deep Learning 文章标签： pytorch 深度学习 python

于 2022-06-29 21:21:55 首次发布

本文链接：https://blog.youkuaiyun.com/xiangyong58/article/details/125529645

Machine & Deep Learning 专栏收录该内容

78 篇文章 ¥9.90 ¥99.00

订阅专栏

超级会员免费看

本文详细介绍了PyTorch中的DataParallel（DP）、DistributedDataParallel（DDP）以及如何进行多机多卡的分布式训练。通过实例解析了如何利用torchrun进行分布式训练的配置，是PyTorch深度学习项目中进行并行计算的重要参考。

1、DataParallel

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
###	第一步：构建模型
'''
model 需要分发的模型
device_ids 可分发的gpu，默认分发到所有看见GPU（环境变量设置的）
output_device 结果输出设备 通常设置成逻辑gpu的第一个
'''

model = nn.DataParallel( model, device_ids=range(args.gpus), output_device=None)

### 第二步：数据迁移
inputs=inputs.to(device)	
labels=labels.to(device)	
#此处的device通常应为模型输出的output_device，否则无法计算loss

2、DistributedDataParallel

#运行方式一：（旧）
python -m torch.distributed.launch \
       --nnodes 1 \
       --nproc_per_node=4 \
       YourScript.py
# nnodes: 表示有多少个节点，可以通俗的理解为有多少台机器
# nproc_per_node 表示每个节点上有多少个进程，每个进程一般独占一块GPU


#运行方式二：
torchrun main.py #除了--use_env参数，其他torch.distributed.launch所使用的参数均可使用 #如nnodes、nproc_per_node

了解本专栏