PyTorch 101，Part4：內存管理以及使用多個GPU
這篇文章涉及到PyTorch高級的GPU管理，包括如何爲你的神經網絡配置多個GPU，是否使用數據或者模型的並行計算。我們用最佳的實踐總結了如何調試內存錯誤。
這裏是我們PyTorch 101系列的第四部分，在這篇文章中，我們將會涉及到多GPU的使用。
在這個部分我們將會涉及：
- 如何讓你的網絡使用多個GPU，使用數據或者模型並行計算。
- 當創建一個新的對象的時候，如何自動地選擇GPU
- 如何診斷和分析出現的內存問題
目錄
PyTorch 101，Part4：內存管理以及使用多個GPU
1 在CPU和GPUs之間移動張量
在PyTorch中，每個張量都有一個to()的成員變量。它的任務是將張量放在一個具體設備上，可能是CPU或者是GPU。to函數的輸入是一個torch.device的實例對象，它可以用下面的輸入進行初始化：
- cpu:將數據放在CPU上
- cuda:0 將數據放在第0個GPU上。同樣，你可以把張量放置在別的GPU上面。
一般來說，無論什麼時候你初始化一個Tensor，它都會默認放在CPU上面。你之後可以將它移動到GPU上面。你可以通過調用torch.cuda.is_available函數來檢查GPU是否可用。
if torch.cuda.is_available():
dev = "cuda:0"
else:
dev = "cpu"
device = torch.device(dev)
# 或者使用
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
a = torch.zeros(4,3)
a = a.to(device) #alternatively, a.to(0)
你也可以通過傳入一個索引作爲to函數的參數，將張量移動到某個具體的GPU上面。重要的是，上面的代碼和設備無關，也就是說，你不需要單獨改變它，就可以讓數據既能在CPU上運行又能在GPU上運行（這完全根據你電腦的配置來定）。
1.1 cuda()函數
另一個將張量放置在GPU上的方法是調用cuda(n)函數，其中n代表的是GPU的索引。如果你調用了cuda，那麼這個張量就會放置在CPU n上面。
torch.nn.Module類還有to函數和cuda函數，它可以將整個網絡放在一個特定的設備上面。Tensor在nn.Module的實例對象上調用to函數就夠了，並不需要分配to函數的返回值。
clf = myNetwork()
clf.to(torch.device("cuda:0") # or clf = clf.cuda()
# 或者使用：clf = myNetwork().to(device)
2 自動選擇GPU
儘管可以顯式地定義張量在哪個GPU上運行，但是在進行運算的過程中，我們會創建很多張量，我們希望它能夠自動地在某個設備上進行創建，這樣能夠減少張量在設備間進行交換，這會降低代碼執行的速度。在這方面，PyTorch給我們提供了一些函數去實現這個思路。
第一個函數是torch.get_device()函數。它只支持GPU類型的張量，這個函數的返回值張量所在GPU的索引。我們可以使用這個函數取決定張量的所在的設備，因此我們可以將一個創建好了的張量自動地移動到這個設備上。
#making sure t2 is on the same device as t2
a = t1.get_device()
b = torch.tensor(a.shape).to(dev)
當我們創建一個新的Tensors的時候，我們也可以調用cuda(n)函數。通過調用cuda函數所創建的張量默認放置在GPU0上面，但是通過下面的聲明對其進行改變。
torch.cuda.set_device(0) # or 1,2,3
如果在兩個操作數之間創建了一個張量，這兩個操作數在相同的設備上，那麼這個張量也會放置在這個設備上。但是如果兩個操作數在不同的設備上上面，這將會報錯。
2.1 new_*函數
在PyTorch1.0版本中，還可以利用new_函數實現這樣的結果。當一個Tensor調用了形如new_ones函數的時候，它會返回一個和這個Tensor相同的數據類型，以及將其放置在調用new_ones張量相同的設備上面。
ones = torch.ones((2,)).cuda(0)
# Create a tensor of ones of size (3,4) on same device as of "ones"
newOnes = ones.new_ones((3,4))
randTensor = torch.randn(2,4)
更多new_函數可也在PyTorch文檔中找到，相關連接我放在下面了。
3 使用多個GPUs
我們有兩個使用多個GPUs的方法：
- 數據並行計算。我們可以將一個大的batches分成很多小一點的batches，並且在多個GPUs上面並行處理這些小的batches
- 模型並行計算。我們可以將一個神經網絡分成幾個小一點的子網絡，然後再不同的GPU上面處理這些子網絡。
3.1 數據並行計算
在PyTorch中中，可以使用nn.DataParallel類來實現數據並行計算。你可以用一個表示你的神經網絡的nn.Module實例對象來初始化一個nn.DataParallel實例對象，然後傳入一連串的GPU索引號，這些表示batches要在哪些GPU上進行處理。
parallel_net = nn.DataParallel(myNet, gpu_ids = [0,1,2])
現在，你可以簡單地執行nn.DataParallel對象，就像使用一個nn.Module對象一樣。
predictions = parallel_net(inputs) # Forward pass on multi-GPUs
loss = loss_function(predictions, labels) # Compute loss function
loss.mean().backward() # Average GPU-losses + backward pass
optimizer.step()
但是，我需要說清楚一些事情。儘管我們的數據在多個GPUs上並行計算，但是在一開始我們需要將數據保存在一個GPU上面。
我們還需要確保DataParallel實例對象也在特定的GPU上面。這個語法和我們之前在使用nn.Module的相似。
input = input.to(0)
parallel_net = parellel_net.to(0)
事實上，下面的示例圖表述了nn.DataParallel是如何工作的。
DataParallel接收輸入，將數據分成多個更小的batches，在所有的GPU上覆制神經網絡，執行前向傳播，然後再原來的GPU上整合輸出。
DataParallel存在一個問題，那就是在一個GPU上面（就是主節點）會不均衡的加載數據。通常來說有兩個方法可以解決這些問題。
- 首先，在前向傳播的時候計算損失。這可以確保至少是並行計算網絡的損失函數。
- 另一個方法是去實現一個並行的損失函數網絡層。這個超過這這篇文章的範圍，但是，如果感興趣的話，我給出了相關文章的鏈接，他們詳細介紹瞭如何實現這樣的網絡層。
3.2 模型並行計算
模型並行計算意味着你可以將你的網絡模型劃分成多個更小的子網絡，然後將他們放在不同的GPU上面。這樣做的主要是因爲如果你的網絡模型太大，一個GPU放不下。
注意，模型並行計算比數據並行計算要慢，這是因爲將一個網絡模型放在多個GPUs上面會增加GPUs之間的以來，這樣會妨礙它們完全地進行並行計算。這樣並不是出於速度快的優點，而是因爲它可以在網絡模型尺寸太大，一個GPU放不下的時候去運行網絡模型。
正如圖b所示，在前向傳播的時候，子網絡2在等待子網絡1運行完成，而在反向傳播的時候，子網絡1要等到子網絡2運行完成。
只要你記住兩件事，在PyTorch實現模型並行計算就會非常的容易。
- 輸入數據和網絡模型應該在相同的設備上面。
- to()函數和cuda()函數都支持autograd，因此在反向傳播的過程中，你可以把梯度從一個GPU複製到另外一個GPU上面。
我們使用下面的代碼來更好的理解。
class model_parallel(nn.Module):
def __init__(self):
super().__init__()
self.sub_network1 = ...
self.sub_network2 = ...
self.sub_network1.cuda(0)
self.sub_network2.cuda(1)
def forward(x):
x = x.cuda(0)
x = self.sub_network1(x)
x = x.cuda(1)
x = self.sub_network2(x)
return x
在init()函數中，我們單獨將子網絡放在GPU0和1上面。
注意，在forward()函數中，在將中間數據傳給sub_network2之前，我們將其從sub_network1轉爲了GPU 1。因爲cuda支持autograd，爲了接下來的反向傳播，sub_network2的損失的反向傳播會複製到sub_network1的緩存器上。
4 解決內存溢出(OOM)錯誤
在這部分中，我們會涉及如何診斷內存問題，以及如果你的網絡超出它需要的內存的可能的解決方法。
儘管內存溢出可能需要減少批次的大小，但是最佳的手段是去檢查內存的使用。
4.1 使用GPUtil去跟蹤內存的使用情況
跟蹤GPU使用的方法是使用nvidia-smi命令行控制檯上監控內存的使用情況。這個方法有個問題就是會泄露GPU的使用情況，並且內存泄露發生的太快，你無法準確定位你的代碼的哪個部分導致了這個錯誤。
爲了解決這個問題，我們可以使用GPUtil擴展模塊，你可以運行下面的命令，使用pip進行安裝。
pip install GPUtil
它的使用也非常的簡單
import GPUtil
GPUtil.showUtilization()
這需要把第二行代碼放在你想檢查GPU使用情況的地方就行。通過將這行代碼放在代碼中的不同位置，你就可以檢查究竟是哪個部分導致了你的網絡模型的OOM錯誤。
讓我們討論一下解決OOM錯誤可能的方法
4.2 使用del關鍵字處理內存丟失
PyTorch有着非常積極的垃圾回收機制。只要變量超出了範圍，垃圾回收就會釋放它。
我們需要記住，Python不會像其他語言比如c/c++那樣強制執行上面的規則。只要一個變量沒有其他指針指向它的時候，它就會被釋放（這也就是變量在Python中不需要聲明的原因）。
事實上，你的input，output張量所佔用的內存儘管超出了訓練範圍，但是仍然不會被釋放。思考一下下面的代碼：
運行上面的代碼，仍會輸出i的值，儘管已
for x in range(10):
i = x
print(i) # 9 is printed
儘管超出了我們初始化i的循環。同樣的，loss和output的張量在訓練循環之外依然存在。我們可以使用del關鍵字來真正地釋放這些張量。
del out, loss
作爲一般的經驗法則，如果你處理了一個張量，你應該del它，因爲它不會被當做垃圾進行回收，除非沒有指針指向它。
4.3 使用Python數據格式而非一維的張量
在訓練循環中，我們通常對一些值進行求和去計算一些矩陣。一個大型的例子就是在每次迭代中，我們需要去更新損失。但是，如果在PyTorch中處理的不小心的話，一些事情可能會導致超出它所需要的內存的使用。
思考下面的代碼：
total_loss = 0
for x in range(10):
# assume loss is computed
iter_loss = torch.randn(3,4).mean()
iter_loss.requires_grad = True # losses are supposed to differentiable
total_loss += iter_loss # use total_loss += iter_loss.item) instead
在接下來的迭代中，我們需要iter_loss的引用會指定到一個新的iter_loss上，然後代表之前聲明的iter_loss對象會被釋放。但是事與願違，爲什麼？
因爲iter_loss是可微的，total_loss += iter_loss這行代碼使用一個AddBakward函數結點構建了一個計算圖。在接下來的迭代中，AddBakward結點會添加到這個圖中，而且保存iter_loss的對象不會被釋放。通常來說，當backward調用的時候，分配給計算圖的內存會被釋放，但是在這裏，無法調用backward函數。
解決這個問題的方法是給total_loss加上一個Python數據，而非一個張量，這可能有效的預防計算圖的創建。
我們僅僅使用total_loss+=iter_loss代替total_loss++iter_loss.item()。tiem從一個值保存數值的張量中返回python數據。
4.4 清空Cuda寄存器
儘管PyTorch會積極地釋放內存，但是PyTorch進程可能不會將內存返回給操作系統，儘管你已經del了你的張量了。這個內存已經保存了，因此它可以很快的分配給一個新的張量，而不請求操作系統新的額外內存。
當你在工作流程中使用了超過兩個進程的時候，將會報錯。
當第二個進行運行的時候，儘管第一個進程導致了OOM，但是它還是會佔用GPU的內存。爲了解決這個問題，你可以拿個在代碼的最後加上這行命令。
torch.cuda.empy_cache()
我們需要保證被進程佔用的空間已經釋放了。
import torch
from GPUtil import showUtilization as gpu_usage
print("Initial GPU Usage")
gpu_usage()
tensorList = []
for x in range(10):
tensorList.append(torch.randn(10000000,10).cuda()) # reduce the size of tensor if you are getting OOM
print("GPU Usage after allcoating a bunch of Tensors")
gpu_usage()
del tensorList
print("GPU Usage after deleting the Tensors")
gpu_usage()
print("GPU Usage after emptying the cache")
torch.cuda.empty_cache()
gpu_usage()
下面的結果是在Tesla K80上執行上面代碼的結果：
Initial GPU Usage
| ID | GPU | MEM |
------------------
| 0 | 0% | 5% |
GPU Usage after allcoating a bunch of Tensors
| ID | GPU | MEM |
------------------
| 0 | 3% | 30% |
GPU Usage after deleting the Tensors
| ID | GPU | MEM |
------------------
| 0 | 3% | 30% |
GPU Usage after emptying the cache
| ID | GPU | MEM |
------------------
| 0 | 3% | 5% |
4.5 在推理中使用torch.no_grad()
PyTorch在前向傳播的時候默認會構建一個計算圖。在創建計算圖的過程中，它會分配緩存去存儲梯度和中間值，在反向傳播的過程中，會使用它們去計算梯度。
在反向傳播的過程中，除了爲葉子結點分配的緩存器之外，所有的緩存器都是被釋放。
但是，在推理的過程中，不存在反向傳播，並且這些緩存器從來不會釋放，這樣會導致內存堆積。因此，當你想去執行的代碼不需要進行反向傳播的時候，將代碼放在torch.no_grad()上下文中。
with torch.no_grad()
# your code
5 總結
在這裏總結了關於內存管理和PyTorch中多個GPUs的使用。下面是一些你可能想去跟進的文章的重要鏈接。
6 擴展閱讀
- PyTorch
newfunctions - Parallelised Loss Layer: Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups
- GPUtil Github page