8.4 多GPU計算
注：相對于本章的前面幾節(jié)，我們實際中更可能遇到本節(jié)所討論的情況：多GPU計算版姑。原書將MXNet的多GPU計算分成了8.4和8.5兩節(jié)，但我們將關于PyTorch的多GPU計算統一放在本節(jié)討論。 需要注意的是味咳，這里我們談論的是單主機多GPU計算而不是分布式計算。如果對分布式計算感興趣可以參考PyTorch官方文檔。
本節(jié)中我們將展示如何使用多塊GPU計算腊尚，例如略荡，使用多塊GPU訓練同一個模型庵佣。正如所期望的那樣，運行本節(jié)中的程序需要至少2塊GPU汛兜。事實上巴粪，一臺機器上安裝多塊GPU很常見，這是因為主板上通常會有多個PCIe插槽。如果正確安裝了NVIDIA驅動肛根，我們可以通過在命令行輸入nvidia-smi命令來查看當前計算機上的全部GPU（或者在jupyter notebook中運行!nvidia-smi）辫塌。
nvidia-smi
輸出：
Wed May 15 23:12:38 2019
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 390.48 Driver Version: 390.48 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
|===============================+======================+======================|
| 0 TITAN X (Pascal) Off | 00000000:02:00.0 Off | N/A |
| 46% 76C P2 87W / 250W | 10995MiB / 12196MiB | 0% Default |
+-------------------------------+----------------------+----------------------+
| 1 TITAN X (Pascal) Off | 00000000:04:00.0 Off | N/A |
| 53% 84C P2 143W / 250W | 11671MiB / 12196MiB | 4% Default |
+-------------------------------+----------------------+----------------------+
| 2 TITAN X (Pascal) Off | 00000000:83:00.0 Off | N/A |
| 62% 87C P2 190W / 250W | 12096MiB / 12196MiB | 100% Default |
+-------------------------------+----------------------+----------------------+
| 3 TITAN X (Pascal) Off | 00000000:84:00.0 Off | N/A |
| 51% 83C P2 255W / 250W | 8144MiB / 12196MiB | 58% Default |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: GPU Memory |
| GPU PID Type Process name Usage |
|=============================================================================|
| 0 44683 C python 3289MiB |
| 0 155760 C python 4345MiB |
| 0 158310 C python 2297MiB |
| 0 172338 C /home/yzs/anaconda3/bin/python 1031MiB |
| 1 139985 C python 11653MiB |
| 2 38630 C python 5547MiB |
| 2 43127 C python 5791MiB |
| 2 156710 C python3 725MiB |
| 3 14444 C python3 1891MiB |
| 3 43407 C python 5841MiB |
| 3 88478 C /home/tangss/.conda/envs/py36/bin/python 379MiB |
+-----------------------------------------------------------------------------+
從上面的輸出可以看到一共有四塊TITAN X GPU，每一塊總共有約12個G的顯存派哲，此時每塊的顯存都占得差不多了......此外還可以看到GPU利用率臼氨、運行的所有程序等信息。
Pytorch在0.4.0及以后的版本中已經提供了多GPU訓練的方式芭届，本文用一個簡單的例子講解下使用Pytorch多GPU訓練的方式以及一些注意的地方储矩。
8.4.1 多GPU計算
先定義一個模型：
import torch
net = torch.nn.Linear(10, 1).cuda()
net
輸出：
Linear(in_features=10, out_features=1, bias=True)
要想使用PyTorch進行多GPU計算，最簡單的方法是直接用torch.nn.DataParallel將模型wrap一下即可：
net = torch.nn.DataParallel(net)
net
輸出：
DataParallel(
(module): Linear(in_features=10, out_features=1, bias=True)
)
這時褂乍，默認所有存在的GPU都會被使用持隧。
如果我們機子中有很多GPU(例如上面顯示我們有4張顯卡，但是只有第0逃片、3塊還剩下一點點顯存)屡拨，但我們只想使用0、3號顯卡褥实，那么我們可以用參數device_ids指定即可:torch.nn.DataParallel(net, device_ids=[0, 3])洁仗。
8.4.2 多GPU模型的保存與加載
我們現在來嘗試一下按照4.5節(jié)（讀取和存儲）推薦的方式進行一下模型的保存與加載。 保存模型:
torch.save(net.state_dict(), "./8.4_model.pt")
加載模型前我們一般要先進行一下模型定義性锭，此時的new_net并沒有使用多GPU:
new_net = torch.nn.Linear(10, 1)
new_net.load_state_dict(torch.load("./8.4_model.pt"))
然后我們發(fā)現報錯了:
RuntimeError: Error(s) in loading state_dict for Linear:
Missing key(s) in state_dict: "weight", "bias".
Unexpected key(s) in state_dict: "module.weight", "module.bias".
事實上DataParallel也是一個nn.Module赠潦，只是這個類其中有一個module就是傳入的實際模型。因此當我們調用DataParallel后草冈，模型結構變了（在外面加了一層而已她奥，從8.4.1節(jié)兩個輸出可以對比看出來）。所以直接加載肯定會報錯的怎棱，因為模型結構對不上哩俭。
所以正確的方法是保存的時候只保存net.module:
torch.save(net.module.state_dict(), "./8.4_model.pt")
new_net.load_state_dict(torch.load("./8.4_model.pt")) # 加載成功
或者先將new_net用DataParallel包括以下再用上面報錯的方法進行模型加載:
torch.save(net.state_dict(), "./8.4_model.pt")
new_net = torch.nn.Linear(10, 1)
new_net = torch.nn.DataParallel(new_net)
new_net.load_state_dict(torch.load("./8.4_model.pt")) # 加載成功
注意這兩種方法的區(qū)別，推薦用第一種方法拳恋，因為可以按照普通的加載方法進行正確加載凡资。
注：本節(jié)與原書基本不同，原書傳送門
