【code】多GPU多卡并行如何操作

最初以为多卡并行仅仅只是一句把模型model以及optimize放在nn.DataParallel上就可以完成：

import torch.nn as nn

【1.返回模型】

self.bert_model = nn.DataParallel(self.bert_model, device_ids=args.gpu_ids, output_device=torch.cuda.current_device())

或【2.返回计算结果】

nn.parallel.data_parallel(module, inputs, device_ids=args.gpu_ids, output_device=torch.cuda.current_device(), dim=0, module_kwargs=None)

但并非是这样。如果单单使用这行命令会发现原本正确的程序开始报错。

正确的使用方法为：

model:

model = nn.DataParallel(model, device_ids=args.bert_gpu_ids) #, output_device=torch.cuda.current_device())

optimizer:

optimizer = torch.optim.Adam([param for _, param in param_optimizer],lr=args.lr)

optimizer = nn.DataParallel(optimizer, device_ids=args.bert_gpu_ids) #, output_device=torch.cuda.current_device())

for param_group in optimizer.module.param_groups:
    param_group['lr'] = cur_lr 

optimizer.module.step() 

 

 

NOTE :

1. 主 GPU 默认情况下是 0 号 GPU，也可以通过 torch.cuda.set_device(id) 来手动更改默认 GPU。
2. 提供的多 GPU 并行列表中需要包含有主 GPU。
3. torch.cuda.set_device(args.device):设置当前设备。官网不鼓励使用此函数来设置。在大多数情况下，最好使用CUDA_VISIBLE_DEVICES环境变量。
    

 

https://pytorch-cn.readthedocs.io/zh/latest/package_references/torch-cuda/

eg:
os.environ['CUDA_VISIBLE_DEVICES'] = "5,6,7"

net = torch.nn.DataParallel(net, device_ids=[0,1,2])

 

参考博客：https://blog.csdn.net/qq_19598705/article/details/80396325