Pytorch訓練時視訊記憶體分配過程探究

頎周發表於2020-12-09
PyTorch記憶體

　　對於視訊記憶體不充足的煉丹研究者來說，弄清楚Pytorch視訊記憶體的分配機制是很有必要的。下面直接通過實驗來推出Pytorch視訊記憶體的分配過程。

　　實驗實驗程式碼如下：

import torch 
from torch import cuda 

x = torch.zeros([3,1024,1024,256],requires_grad=True,device='cuda') 
print("1", cuda.memory_allocated()/1024**2)  
y = 5 * x 
print("2", cuda.memory_allocated()/1024**2)  
torch.mean(y).backward()     
print("3", cuda.memory_allocated()/1024**2)    
print(cuda.memory_summary())

　　輸出如下：

　　程式碼首先分配3GB的視訊記憶體建立變數x，然後計算y，再用y進行反向傳播。可以看到，建立x後與計算y後分別佔視訊記憶體3GB與6GB，這是合理的。另外，後面通過backward()，計算出x.grad，佔存與x一致，所以最終一共佔有視訊記憶體9GB，這也是合理的。但是，輸出顯示了視訊記憶體的峰值為12GB，這多出的3GB是怎麼來的呢？首先畫出計算圖：

　　下面通過列表的形式來模擬Pytorch在運算時分配視訊記憶體的過程：

 

　　如上所示，由於需要儲存反向傳播以前所有前向傳播的中間變數，所以有了12GB的峰值佔存。

　　我們可以不儲存計算圖中的非葉子結點，達到節省視訊記憶體的目的，即可以把上面的程式碼中的y=5*x與mean(y)寫成一步：

import torch 
from torch import cuda 

x = torch.zeros([3,1024,1024,256],requires_grad=True,device='cuda') 
print("1", cuda.memory_allocated()/1024**2)    
torch.mean(5*x).backward()     
print("2", cuda.memory_allocated()/1024**2)    
print(cuda.memory_summary())

　　佔視訊記憶體量減少了3GB：

相關文章