Pytorch顯示記憶體動態分配規律探索

• 2020 年 11 月 16 日
• 筆記
• TF/Keras/Pytorch

　　下面通過實驗來探索Pytorch分配顯示記憶體的方式。

實驗

顯示記憶體到主存

　　我使用VSCode的jupyter來進行實驗，首先只導入pytorch，程式碼如下：

import torch

　　打開任務管理器查看主存與顯示記憶體情況。情況分別如下：

　　在顯示記憶體中創建1GB的張量，賦值給a，程式碼如下：

a = torch.zeros([256,1024,1024],device= 'cpu') 

　　查看主存與顯示記憶體情況：

　　可以看到主存與顯示記憶體都變大了，而且顯示記憶體不止變大了1G，多出來的記憶體是pytorch運行所需的一些配置變數，我們這裡忽略。

　　再次在顯示記憶體中創建一個1GB的張量，賦值給b，程式碼如下：

b = torch.zeros([256,1024,1024],device= 'cpu') 

　　查看主顯示記憶體情況：

　　這次主存大小沒變，顯示記憶體變高了1GB，這是合情合理的。然後我們將b移動到主存中，程式碼如下：

b = b.to('cpu')  

　　查看主顯示記憶體情況：

　　發現主存是變高了1GB，顯示記憶體卻只變小了0.1GB，好像只是將顯示記憶體張量複製到主存一樣。實際上，pytorch的確是複製了一份張量到主存中，但它也對顯示記憶體中這個張量的移動進行了記錄。我們接著執行以下程式碼，再創建1GB的張量賦值給c：

c = torch.zeros([256,1024,1024],device= 'cuda')  

　　查看主顯示記憶體情況：

　　發現只有顯示記憶體大小變大了0.1GB，這說明，Pytorch的確記錄了顯示記憶體中張量的移動，只是沒有立即將顯示記憶體空間釋放，它選擇在下一次創建新變數時覆蓋這個位置。接下來，我們重複執行上面這行程式碼：

c = torch.zeros([256,1024,1024],device= 'cuda')  

　　主顯示記憶體情況如下：

　　明明我們把張量c給覆蓋了，顯示記憶體內容卻變大了，這是為什麼呢？實際上，Pytorch在執行這句程式碼時，是首先找到可使用的顯示記憶體位置，創建這1GB的張量，然後再賦值給c。但因為在新創建這個張量時，原本的c依然佔有1GB的顯示記憶體，pytorch只能先調取另外1GB顯示記憶體來創建這個張量，再將這個張量賦值給c。這樣一來，原本的那個c所在的顯示記憶體內容就空出來了，但和前面說的一樣，pytorch並不會立即釋放這裡的顯示記憶體，而等待下一次的覆蓋，所以顯示記憶體大小並沒有減小。

　　我們再創建1GB的d張量，就可以驗證上面的猜想，程式碼如下：

d = torch.zeros([256,1024,1024],device= 'cuda')  

　　主顯示記憶體情況如下：

　　顯示記憶體大小並沒有變，就是因為pytorch將新的張量創建在了上一步c空出來的位置，然後再賦值給了d。另外，刪除變數操作也同樣不會立即釋放顯示記憶體：

del d

　　主顯示記憶體情況：

　　顯示記憶體沒有變化，同樣是等待下一次的覆蓋。

主存到顯示記憶體

　　接著上面的實驗，我們創建直接在主存創建1GB的張量並賦值給e，程式碼如下：

e = torch.zeros([256,1024,1024],device= 'cpu')  

　　主顯示記憶體情況如下：

　　主存變大1GB，合情合理。然後將e移動到顯示記憶體，程式碼如下：

e = e.to('cuda')

　　主顯示記憶體情況如下：

　　主存變小1GB，顯示記憶體沒變是因為上面張量d被刪除沒有被覆蓋，合情合理。說明主存的釋放是立即執行的。

總結 

　　通過上面的實驗，我們了解到，pytorch不會立即釋放顯示記憶體中失效變數的記憶體，它會以覆蓋的方式利用顯示記憶體中的可用空間。另外，如果要重置顯示記憶體中的某個規模較大的張量，最好先將它移動到主存中，或是直接刪除，再創建新值，否則就需要兩倍的記憶體來實現這個操作，就有可能出現顯示記憶體不夠用的情況。 

　　實驗程式碼匯總如下：

#%% 
import torch
#%%
a = torch.zeros([256,1024,1024],device= 'cuda')  
#%%
b = torch.zeros([256,1024,1024],device= 'cuda')  
#%%
b = b.to('cpu')
#%%
c = torch.zeros([256,1024,1024],device= 'cuda')  
#%%
c = torch.zeros([256,1024,1024],device= 'cuda')  
#%%  
d = torch.zeros([256,1024,1024],device= 'cuda')  
#%%
del d 
#%%  
e = torch.zeros([256,1024,1024],device= 'cpu')  
#%%
e = e.to('cuda')