ICode9

精准搜索请尝试: 精确搜索
首页 > 其他分享> 文章详细

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

2022-02-28 19:00:33  阅读:480  来源: 互联网

标签:loss torch args DDP 多卡 rank pytorch sampler local


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

本站声明: 1. iCode9 技术分享网(下文简称本站)提供的所有内容,仅供技术学习、探讨和分享;
2. 关于本站的所有留言、评论、转载及引用,纯属内容发起人的个人观点,与本站观点和立场无关;
3. 关于本站的所有言论和文字,纯属内容发起人的个人观点,与本站观点和立场无关;
4. 本站文章均是网友提供,不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属;如您发现该文章侵犯了您的权益,可联系我们第一时间进行删除;
5. 本站为非盈利性的个人网站,所有内容不会用来进行牟利,也不会利用任何形式的广告来间接获益,纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。

专注分享技术,共同学习,共同进步。侵权联系[81616952@qq.com]

Copyright (C)ICode9.com, All Rights Reserved.

ICode9版权所有