1. 利用CUDA_VISIBLE_DEVICES设置可用显卡

在CUDA中设定可用显卡，一般有2种方式：

(1) 在代码中直接指定

import os 
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_ids 

(2) 在命令行中执行代码时指定

CUDA_VISIBLE_DEVICES=gpu_ids python3 train.py

如果使用sh脚本文件运行代码，则有3种方式可以设置

(3) 在命令行中执行脚本文件时指定：

CUDA_VISIBLE_DEVICES=gpu_ids sh run.sh

(4) 在sh脚本中指定：

source bashrc
export CUDA_VISIBLE_DEVICES=gpu_ids && python3 train.py

(5) 在sh脚本中指定

source bashrc
CUDA_VISIBLE_DEVICES=gpu_ids python3 train.py

如果同时使用多个设定可用显卡的指令，比如

source bashrc
export CUDA_VISIBLE_DEVICES=gpu_id1 && CUDA_VISIBLE_DEVICES=gpu_id2 python3 train.py

那么高优先级的指令会覆盖第优先级的指令使其失效。

优先级顺序为：不使用sh脚本 (1)>(2)； 使用sh脚本(1)>(5)>(4)>(3)

个人感觉在炼丹时建议大家从(2)(3)(4)(5)中选择一个指定可用显卡，不要重复指定以防造成代码的混乱。方法(1)虽然优先级最高，但是需要修改源代码，所以不建议使用。

 

2 .cuda()方法和torch.cuda.set_device()

可以使用.cuda()[包括model.cuda()/loss.cuda()/tensor.cuda()]方法和torch.cuda.set_device()来把模型和数据加载到对应的gpu上。

(1) .cuda()

以model.cuda()为例，加载方法为：

model.cuda(gpu_id) # gpu_id为int类型变量，只能指定一张显卡
model.cuda('cuda:'+str(gpu_ids)) #输入参数为str类型，可指定多张显卡
model.cuda('cuda:1,2') #指定多张显卡的一个示例

(2) torch.cuda.set_device()

使用torch.cuda.set_device()可以更方便地将模型和数据加载到对应GPU上, 直接定义模型之前加入一行代码即可

torch.cuda.set_device(gpu_id) #单卡
torch.cuda.set_device('cuda:'+str(gpu_ids)) #可指定多卡

但是这种写法的优先级低，如果model.cuda()中指定了参数，那么torch.cuda.set_device()会失效，而且pytorch的官方文档中明确说明，不建议用户使用该方法。

第1节和第2节所说的方法同时使用是并不会冲突，而是会叠加。比如在运行代码时使用

CUDA_VISIBLE_DEVICES=2,3,4,5 python3 train.py

而在代码内部又指定

model.cuda(1)
loss.cuda(1)
tensor.cuda(1)

那么代码会在GPU3上运行。原理是CUDA_VISIBLE_DEVICES使得只有GPU2,3,4,5可见，那么这4张显卡，程序就会把它们看成GPU0,1,2,3，.cuda(1)把模型/loss/数据都加载到了程序所以为的GPU1上，则实际使用的显卡是GPU3。

如果利用.cuda()或torch.cuda.set_device()把模型加载到多个显卡上，而实际上只使用一张显卡运行程序的话，那么程序会把模型加载到第一个显卡上，比如如果在代码中指定了

model.cuda('cuda:2,1')

在运行代码时使用

CUDA_VISIBLE_DEVICES=2,3,4,5 python3 train.py

这一指令，那么程序最终会在GPU4上运行。

 

3.多卡数据并行torch.nn.DataParallel

多卡数据并行一般使用

torch.nn.DataParallel(model,device_ids)

其中model是需要运行的模型，device_ids指定部署模型的显卡，数据类型是list

device_ids中的第一个GPU（即device_ids[0]）和model.cuda()或torch.cuda.set_device()中的第一个GPU序号应保持一致，否则会报错。此外如果两者的第一个GPU序号都不是0,比如设置为：

model=torch.nn.DataParallel(model,device_ids=[2,3])
model.cuda(2)

那么程序可以在GPU2和GPU3上正常运行，但是还会占用GPU0的一部分显存（大约500M左右），这是由于pytorch本身的bug导致的（截止1.4.0，没有修复这个bug）。

device_ids的默认值是使用可见的GPU，不设置model.cuda()或torch.cuda.set_device()等效于设置了model.cuda(0)

4. 多卡多线程并行torch.nn.parallel.DistributedDataParallel

（这个我是真的没有搞懂,,,,）

参考了这篇文章和这个代码，关于GPU的指定，多卡多线程中有2个地方需要设置

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

模型/loss/tensor设置为.cuda()或.cuda(args.local_rank)均可，不影响正常运行。

 

5. 推荐设置方式：

(1) 单卡

使用CUDA_VISIBLE_DEVICES指定GPU，不要使用torch.cuda.set_device()，不要给.cuda()赋值。

(2) 多卡数据并行

直接指定CUDA_VISIBLE_DEVICES，通过调整可见显卡的顺序指定加载模型对应的GPU,不要使用torch.cuda.set_device()，不要给.cuda()赋值，不要给torch.nn.DataParallel中的device_ids赋值。比如想在GPU1,2,3中运行，其中GPU2是存放模型的显卡，那么直接设置

CUDA_VISIBLE_DEVICES=2,1,3

(3) 多卡多线程

推荐参考这篇文章和这个代码

 

参考

当代研究生应当掌握的并行训练方法（单机多卡） - 知乎 (zhihu.com) 

pytorch-distributed/distributed.py at master · tczhangzhi/pytorch-distributed · GitHub