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

1、DataParallel

1.  
   device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
2.  
   ### 第一步：构建模型
3.  
   '''
4.  
   model 需要分发的模型
5.  
   device_ids 可分发的gpu，默认分发到所有看见GPU（环境变量设置的）
6.  
   output_device 结果输出设备 通常设置成逻辑gpu的第一个
7.  
   '''
8.  
    
9.  
   model = nn.DataParallel( model, device_ids=range(args.gpus), output_device=None)
10.  
     
11.  
    ### 第二步：数据迁移
12.  
    inputs=inputs.to(device)
13.  
    labels=labels.to(device)
14.  
    #此处的device通常应为模型输出的output_device，否则无法计算loss

2、DistributedDataParallel

1.  
   #运行方式一：（旧）
2.  
   python -m torch.distributed.launch \
3.  
   --nnodes 1 \
4.  
   --nproc_per_node=4 \
5.  
   YourScript.py
6.  
   # nnodes: 表示有多少个节点，可以通俗的理解为有多少台机器
7.  
   # nproc_per_node 表示每个节点上有多少个进程，每个进程一般独占一块GPU
8.  
    
9.  
    
10.  
    #运行方式二：
11.  
    torchrun main.py #除了--use_env参数，其他torch.distributed.launch所使用的参数均可使用 #如nnodes、nproc_per_node
12.  
     
13.  
    OMP_NUM_THREADS=1 torchrun --nnodes

这篇好文章是转载于：学新通技术网