PyTorch指定GPU并行训练小记_torch.device('cuda') 多gpu

假设一台机器上有多块GPU。给定需要训练的模型，每块GPU将分别独立维护一份完整的模型参数。在模型训练的任意一次迭代中，给定一个小批量，将该批量中的样本划分成若干份并分给每块GPU一份。然后，每块GPU将分别根据自己分到的训练数据样本和自己维护的模型参数计算模型参数的梯度。接下来把这若干块GPU上分别计算得到的梯度相加，从而得到当前的小批量梯度。之后，每块GPU都使用这个小批量梯度分别更新自己维护的那一份完整的模型参数。

在多GPU的机器上，要选择特定的GPU来运行程序，可在程序运行命令前使用命令：

CUDA_VISIBLE_DEVICES=0
1

其中0为服务器中的GPU编号，例如在有8块GPU的机器上可以为0，1，2，3，4，5，6，7等，编号可以通过以下命令查看：

nvidia-smi
1

通过指定GPU编号表明对当前启动多程序或者脚本可见的GPU编号，未指定的GPU则不会被该程序或脚本使用。

需要注意，由CUDA_VISIBLE_DEVICES指定的GPU编号和代码中索引的GPU编号不完全对等，具体可见下面的例子：

CUDA_VISIBLE_DEVICES=1     # 只有编号为1的GPU对程序是可见的，在代码中gpu[0]指的就是这块儿GPU
CUDA_VISIBLE_DEVICES=0,2,3 # 只有编号为0,2,3的GPU对程序是可见的，在代码中gpu[0]指的是第0块儿，
						   # gpu[1]指的是第2块儿，gpu[2]指的是第3块儿
CUDA_VISIBLE_DEVICES=2,0,3 # 只有编号为0,2,3的GPU对程序是可见的，但是在代码中gpu[0]指的是第2块儿，
						   # gpu[1]指的是第0块儿，gpu[2]指的是第3块儿
1
2
3
4
5

torch.device(f"cuda:{args.local_rank}")#args.local_rank由PyTorch自己根据GPU使用情况指定

torch.device("cuda:0")#gpu[0]

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
1
2
3
4
5

CUDA_VISIBLE_DEVICES=0,1,2,3 python gpu_test.py
1

---

参考资料：
[1] 『MXNet』第七弹_多GPU并行程序设计
[2] CUDA GPU 隔离 CUDA_VISIBLE_DEVICES