參考資料：
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的那張卡。