pytorch DistributedDataParallel多卡并行训练

pytorch DistributedDataParallel多卡并行训练

Pytorch 中最简单的并行计算方式是 nn.DataParallel。
DataParallel 使用单进程控制将模型和数据加载到多个 GPU 中,控制数据在 GPU 之间的流动,协同不同 GPU 上的模型进行并行训练。

但是DataParallel的缺点十分明显,各卡之间的负载不均衡,主卡的负载过大。运行时间大概是distributed 的四倍。
所以,下面我们介绍使用distributed 的计算方式。

先设定下local_rank,这是很重要的参数

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=0, type=int, help='node rank for distributed training')
args = parser.parse_args()
print(args.local_rank)
 torch.distributed.init_process_group('nccl', init_method='env://')

 torch.cuda.set_device(args.local_rank)
 device = torch.device("cuda", args.local_rank)

 train_dataset = ModelNet_Loader(root, 'train', args)  #自己写的类
 test_dataset = ModelNet_Loader(root, 'test', args)    #自己写的类

 train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
 test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset)

 train_loader = DataLoader(train_dataset, num_workers=2,batch_size=args.batch_size, shuffle=False, drop_last=True,  sampler=train_sampler)       
 test_loader = DataLoader(test_dataset, num_workers=2,batch_size=args.test_batch_size, shuffle=False, drop_last=False, sampler = test_sampler)
   
    #Try to load models
 if args.model == 'STN3d':
        model = STN3d()    #自己的model
 print(str(model))
 model.to(device)

 model = torch.nn.parallel.DistributedDataParallel(model,device_ids=[args.local_rank],output_device=args.local_rank)

然后我们运行如下代码,启动训练

CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2  --master_port 29501  main.py

参考

https://zhuanlan.zhihu.com/p/98535650

  • 1
    点赞
  • 10
    收藏
    觉得还不错? 一键收藏
  • 5
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 5
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值