1. 程式人生 > >Pytorch視訊記憶體動態分配規律探索

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')