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()


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

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=False)
## 加broadcast_buffers会让loss同步，速度会稍微变慢

3.数据上：

​
## 对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)

​

4. loss上：

保证不同GPU之间的loss是同步的

在loss.backward()后面:

loss.backward()
world_size = dist.get_world_size()
with torch.no_grad():
    dist.all_reduce(loss, async_op=True)
	loss /= world_size

5:运行代码

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

测试结果：

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

（2）精度基本上无变化

标签：loss,torch,args,DDP,多卡,rank,pytorch,sampler,local
来源： https://blog.csdn.net/qq_45009333/article/details/123182276