Pytorch显卡设置
参考资料：
https://zhuanlan.zhihu.com/p/166161217
本人最常使用到显卡和CUDA的东西莫过于Pytorch了。这篇文章着重说明两个问题：1. 如何import torch并使之输出比较完备的CUDA信息 2. 在服务器上有多张卡的环境下，如何使任务在特定的卡或特定的几张卡上跑。
第一个问题：
任务目标是输出信息，那么不妨借助Pytorch的官方示例看一看Pytorch都能输出CUDA的哪些信息。
import torch
from torch.backends import cudnn
x = torch.Tensor([1.0])
xx = x.cuda()
print("Support cudnn ?: ",cudnn.is_acceptable(xx)) # 有些多余
print("Support CUDA ?: ", torch.cuda.is_available())
print("Support {} devices".format(torch.cuda.device_count()))
需要注意的是，这里的device_count是对Pytorch可见的GPU数，细节将在第二个问题中讨论。这里我想介绍一个概念：
这个是我对Pytorch调用CUDA方式的理解，我认为是对的：
1. 默认的Pytorch是仅使用一个GPU的，没错，就算你的总线上有8个GPU，Pytorch默认也只会调用对于它可见的第一个GPU。什么叫仅用一个GPU？即model在这个GPU上被装载，数据也仅被装载到这个GPU上，运算自然也在这个GPU上运行。不信的话可以运行torch.cuda.current_device()，返回的int大部分时候是0，即对于Pytorch可见的第一个GPU。
2. 那么Pytorch是怎么利用多个GPU的？这个问题设计到Pytorch的并行化手段：
DP（DATA PARALLELISM），DDP（DISTRIBUTED DATA PARALLEL），RPC等。想要详细了解的读者可以去看一下官方的tutrial，讲的非常好。
拿最简单的DP并行化手段来说，它利用了单process多thread，涉及了Data上的Parallel，而不涉及model上的Parallel。我做过了一个小实验，现在来描述一下DP并行化发生了什么：
首先，你的（并行化）model被装载到你当前的GPU上，也就是torch.cuda.current_device()返回的GPU上。
然后，你的（并行化）model被复制到其他对Pytorch的GPU上，也就可能是GPU index = 1,2,3...的显卡
再然后，不同的数据被喂给不同的GPU，然后各自独立进行前向和反向。反向传播阶段，不同卡上module的梯度被加和到最初的那张卡的module上。英文版：
This container parallelizes the application of the given module by splitting the input across the specified devices by chunking in the batch dimension (other objects will be copied once per device). In the forward pass, the module is replicated on each device, and each replica handles a portion of the input. During the backwards pass, gradients from each replica are summed into the original module.
第二个问题：
参考资料中将问题说的非常清楚明晰，他同时也给出了推荐的做法。简而言之，合理配置CUDA_VISIBLE_DEVICES这个环境变量是关键。你可以在运行python脚本时（或之前）配置环境变量：
export CUDA_VISIBLE_DEVICES=2,1,0
也可以在python脚本中使用os模块来配置环境变量：
import os os.environ['CUDA_VISIBLE_DEVICES'] = "2,1,0"
第一个问题中我说过，系统中的device序号和Pytorch可见的device序号是不一样的，比如像上面的两端代码设置的一样，Pytorch默认会将nvidia-smi显示的第2,1,0张显卡配置成index = 0, 1, 2。再加上Pytorch默认会在第一个显卡中加载模型和数据，所以如果不进行任何并行化设置的话，一切工作就落在了Pytorch_device_index = 0，CUDA_device_index = 2的那张显卡上！你仍然可以通过torch.cuda.set_device来切换使用哪张卡！比如torch.cuda.set_device(1)，就会使用系统中的CUDA_device_index = 1的那张卡。
