方法1：
理想情况下，该环境变量应设置在程序的顶部。如果在设置 torch.backends.cudnn.benchmark 之后调用 CUDA_VISIBLE_DEVICES 变量，则更改 CUDA_VISIBLE_DEVICES 变量将不起作用。
import os
os.environ["CUDA_VISIBLE_DEVICES"]=<YOUR_GPU_NUMBER_HERE>
方法2：
shell
$ CUDA_VISIBLE_DEVICES=1,2 python train.pyREF
问题：单机多卡使用时，指定device不是从0开始时，出现AssertionError: Invalid device id错误
原因：
pytorch默认使用gpu编号为device:0的设备，可以使用
torch.nn.DataParallel(model, device_ids=[0, 1])
对模型进行制定gpu指定编号多gpu训练，必须要有编号为device:0的gpu，不然会报AssertionError: Invalid device id错误；
如果当gpu编号为device:0的设备被占用时，指定其他编号gpu使用torch.nn.DataParallel(model, device_ids=[1, 2])指定gpu编号会出现AssertionError: Invalid device id错误，这是因为pytorch默认使用gpu编号为device:0的设备，需要修改pytorch默认编号；
解决方法：
方法1：
python
# 使用set_device函数来指定默认GPU使用编号
torch.cuda.set_device(6)
#使用方式
torch.nn.DataParallel(model,device_ids=[6,7,8])缺点：
- 必须包含set_device(1)指定的device:1的设备；
- 且仍然会存在占用一些device:0的gpu内存；
方法2：（推荐）
python
# 使用如下方式来指定使用设备，这样会会把device:6改为device:0,
# device:7改为device:1,以此类推. 则pytorch默认的编号还是以device:0开始.
os.environ["CUDA_VISIBLE_DEVICES"] = "6,7,8"
# 指定方式
torch.nn.DataParallel(model,device_ids=[0,1,2])REF
**问题：**单机多卡训练时，代码只运行了gpu0，并导致CUDA out of memory. 如下图所示：
RuntimeError: CUDA out of memory. Tried to allocate 5.18 GiB (GPU 0; 31.74 GiB total capacity; 21.11 GiB already allocated; 5.10 GiB free; 25.60 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.
解决方法：
with torch.cuda.device(7):
opt = opts().parse()
main(opt)
pytorch在初始化的时候会默认在第0块显卡上进行，会占用一定的显存，这就导致，在第0块显卡空闲内存不多时，程序会反复报上面的错误。
torch.nn.DataParallel
pytorch单机多卡最简单的实现方法就是使用nn.DataParallel类，其几乎仅使用一行代码net = torch.nn.DataParallel(net)就可让模型同时在多张GPU上训练，它大致的工作过程如下图所示：
DataParallel前传与反传工作过程
在每一个Iteration的Forward过程中，nn.DataParallel都自动将输入按照gpu_batch进行split，然后复制模型参数到各个GPU上，分别进行前传后将得到网络输出，最后将结果concat到一起送往0号卡中。
**在Backward过程中，先由0号卡计算loss函数，**通过loss.backward()得到损失函数相于各个gpu输出结果的梯度grad_l1 ... gradln，接下来0号卡将所有的grad_l送回对应的GPU中，然后GPU们分别进行backward得到各个GPU上面的模型参数梯度值gradm1 ... gradmn，最后所有参数的梯度汇总到GPU0卡进行update。
注：DataParallel的整个并行训练过程利用python多线程实现
由以上工作过程分析可知，nn.DataParallel有着这样几个无法避免的问题：
- 负载不均衡问题。gpu0所承担的任务明显要重于其他gpu
- 速度问题。每个iteration都需要复制模型且均从GPU0卡向其他GPU复制，通讯任务重且效率低；python多线程GIL锁导致的线程颠簸(thrashing)问题。
- 只能单机运行。由于单进程的约束导致。
- 只能切分batch到多GPU，而无法让一个model分布在多个GPU上。当一个模型过大，设置batchsize=1时其显存占用仍然大于单张显卡显存，此时就无法使用DataParallel类进行训练。
因此官方推荐使用torch.nn.DistributedDataParallel替代nn.DataParallel，即使是单机多卡的情况下。