pytorch模型加DDP进行单机多卡分布式训练

1.接收local_rank的参数

不要自己替换--local_rank的数值

如果有import导入dataloader，init的代码必须要在dataloader之前。

import argparse
# 运行时，torch.distributed.lunch 会自动传入参数0，1，2来表示是第几个进程
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', type=int, default=-1, help="DDP parameter, do not modify")
args = parser.parse_args()

## 上面代码也可替换成
local_rank = int(os.environ['LOCAL_RANK'])
world_size = int(os.environ['WORLD_SIZE'])
rank = int(os.environ['RANK'])



## Initialization
torch.distributed.init_process_group(backend="nccl" )
torch.cuda.set_device(args.local_rank)

1.1 设置随机数种子，使模型初始化时都是一样的参数

torch.manual_seed(0)
torch.cuda.manual_seed_all(0)

2.模型上：

## 套DDP的时候，model要加.cuda()
model = torch.nn.parallel.DistributedDataParallel(model.cuda(), device_ids=[args.local_rank], output_device=args.local_rank,find_unused_parameters=True)
## 加broadcast_buffers会让loss同步，速度会稍微变慢

3.DataLoader导入数据上：

​
## 对DataLoader中的database进行一个sampler，再在最后加一个sampler=train_sampler即可
train_sampler = torch.utils.data.distributed.DistributedSampler(traindatabase)
trainLoader = torch.utils.data.DataLoader(traindatabase, batch_size=args.batch_size, shuffle = (train_sampler is None), sampler=train_sampler, pin_memory=False)

​
## pin_memory 为 True时，允许在训练的时候提前加载数据，理论上可以提升GPU利用率

3.1 每个epoch中：

for epoch in range(30):
    train_sampler.set_epoch(epoch)
    train(model,xxxxxx)
    xxxxxxx

4. loss上：

默认loss.backward()会等待其他进程的同步。 实验中当设定local_rank == 1的时候sleep，就能看到其他进程不打印了

下面代码能强制保证不同GPU之间的loss是同步的（！！）

在loss.backward()后面:

loss.backward()
## 等待所有进程到这个位置
dist.barrier()   # import torch.distributed as dist
  


## 或者使用：  （还可以计算平均loss）
loss.backward()
world_size = dist.get_world_size()
with torch.no_grad():
    d = dist.all_reduce(loss, async_op=True)
    d.wait()
	loss /= world_size

5:运行代码

## 改成nproc_per_node的值改成GPU数量
python -m torch.distributed.launch --nproc_per_node=2 main.py

6.要使用特定的GPU，在运行的最前面加

CUDA_VISIBLE_DEVICES=3,4 python -m xxxxxxxx

或者在python内开头加

os.environ['CUDA_VISIBLE_DEVICES'] = '3,4'

测试结果：

（1）运行速度： 提升进两倍

（2）精度基本上无变化

注意事项：

        1.在model之外，在任何地方都是默认等待其他进程同步的，

           而在model内的函数操作是不会等待的。（可以利用barrier()或者.wait()）

2. 输入model的data和labels，在dataloader之后在CPU上，一输入model居然自己变成了对应的GPU上

参考：

Pytorch并行分布式训练四 （DDP）_Dteam_f的博客-CSDN博客

pytorch分布式系列3——分布式训练时，torch.utils.data.distributed.DistributedSampler做了什么？_searobbers_duck的博客-CSDN博客