pytorch並行處理詳解

• PyTorch默認使用從0開始的GPU，且默認只使用0號GPU。如果要使用其他編號的GPU或者使用多塊GPU，則要設置。

• pytorch並行後，假設batchsize設置爲64，表示每張並行使用的GPU都使用batchsize=64來計算（單張卡使用時，使用batchsize=64比較合適時，多張卡並行時，batchsize仍爲64比較合適，而不是64*並行卡數）。

• DataParallel 會自動拆分數據，並將作業訂單發送到多個GPU上的多個模型。 在每個模型完成它們的工作之後，DataParallel 在將結果返回給你之前收集和合並結果。

有兩種方法：

方法一：環境變量 + device + to(device)

第一步，指定備選的GPU

有如下兩種方法來指定需要備選的GPU。
這些GPU 將在程序中可見（但不一定會使用，真正投入使用需要通過device()函數來再次選擇和指定）

使用環境變量CUDA_VISIBLE_DEVICES的方式，類似於tensorflow指定GPU的方式（http://www.cnblogs.com/darkknightzh/p/6591923.html）。

直接終端中設定：

CUDA_VISIBLE_DEVICES=1 

python代碼中設定：

1、使用單卡

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

2、使用多塊卡的方法。
例如，使用0號和1號GPU’

import os
os.environ["CUDA_VISIBLE_DEVICES"] = '0,1'

第二步，創建設備（device）

作用：將備選GPU進一步選擇和指定，真正投入使用中。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 注意如果 device = torch.device("cuda")，則環境變量CUDA_VISIBLE_DEVICES中指定的全部GPU都會被拿來使用。
# 也可以通過 "cuda:0" 、"cuda:1"等指定環境變量CUDA_VISIBLE_DEVICES中指定的多塊GPU中的某一塊。

注意對應關係。例如：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3, 4, 5"  # 將2, 3, 4, 5號GPU作爲備選GPU

# 這樣cuda:0表示使用 2 號GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

第三步，將data和model放置到device上

# 如果他們已經在目標設備上則不會執行復制操作
input = data.to(device)
model = MyModule(...).to(device)

注意：如果有多個GPU，則model還需要多做一個操作（模型並行化）
第三步驟的多GPU版本如下：

input = data.to(device)
model = MyModule(...)
#模型並行化
if torch.cuda.device_count() > 1:
            print("Let's use", torch.cuda.device_count(), "GPUs!")
            model = nn.DataParallel(model)
model = model.to(device)

方法二 函數 set_device + 函數.cuda()

不過官方建議使用CUDA_VISIBLE_DEVICES，不建議使用 set_device 函數。

第一步，函數set_device設置device

import torch
gpu_id = [0, 1, 2]
torch.cuda.set_device(gpu_id)

第二部，函數.cuda()使用GPU

data.cuda()
model.cuda()