下面通過實驗來探索Pytorch分配顯存的方式。

實驗顯存到主存

　　我使用VSCode的jupyter來進行實驗，首先只導(dǎo)入pytorch，代碼如下：

import torch

　　打開任務(wù)管理器查看主存與顯存情況。情況分別如下：

　　在顯存中創(chuàng)建1GB的張量，賦值給a，代碼如下：

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

　　查看主存與顯存情況：

　　可以看到主存與顯存都變大了，而且顯存不止變大了1G，多出來的內(nèi)存是pytorch運行所需的一些配置變量，我們這里忽略。

　　再次在顯存中創(chuàng)建一個1GB的張量，賦值給b，代碼如下：

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

　　查看主顯存情況：

　　這次主存大小沒變，顯存變高了1GB，這是合情合理的。然后我們將b移動到主存中，代碼如下：

b = b.to('cpu') 

　　查看主顯存情況：

　　發(fā)現(xiàn)主存是變高了1GB，顯存卻只變小了0.1GB，好像只是將顯存張量復(fù)制到主存一樣。實際上，pytorch的確是復(fù)制了一份張量到主存中，但它也對顯存中這個張量的移動進行了記錄。我們接著執(zhí)行以下代碼，再創(chuàng)建1GB的張量賦值給c：

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

　　查看主顯存情況：

　　發(fā)現(xiàn)只有顯存大小變大了0.1GB，這說明，Pytorch的確記錄了顯存中張量的移動，只是沒有立即將顯存空間釋放，它選擇在下一次創(chuàng)建新變量時覆蓋這個位置。接下來，我們重復(fù)執(zhí)行上面這行代碼：

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

　　主顯存情況如下：

　　明明我們把張量c給覆蓋了，顯存內(nèi)容卻變大了，這是為什么呢？實際上，Pytorch在執(zhí)行這句代碼時，是首先找到可使用的顯存位置，創(chuàng)建這1GB的張量，然后再賦值給c。但因為在新創(chuàng)建這個張量時，原本的c依然占有1GB的顯存，pytorch只能先調(diào)取另外1GB顯存來創(chuàng)建這個張量，再將這個張量賦值給c。這樣一來，原本的那個c所在的顯存內(nèi)容就空出來了，但和前面說的一樣，pytorch并不會立即釋放這里的顯存，而等待下一次的覆蓋，所以顯存大小并沒有減小。

　　我們再創(chuàng)建1GB的d張量，就可以驗證上面的猜想，代碼如下：

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

　　主顯存情況如下：

　　顯存大小并沒有變，就是因為pytorch將新的張量創(chuàng)建在了上一步c空出來的位置，然后再賦值給了d。另外，刪除變量操作也同樣不會立即釋放顯存：

del d

　　主顯存情況：

　　顯存沒有變化，同樣是等待下一次的覆蓋。

主存到顯存

　　接著上面的實驗，我們創(chuàng)建直接在主存創(chuàng)建1GB的張量并賦值給e，代碼如下：

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

　　主顯存情況如下：

　　主存變大1GB，合情合理。然后將e移動到顯存，代碼如下：

e = e.to('cuda')

　　主顯存情況如下：

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

總結(jié)

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

　　實驗代碼匯總?cè)缦拢?/p>

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

到此這篇關(guān)于Pytorch顯存動態(tài)分配規(guī)律探索的文章就介紹到這了,更多相關(guān)Pytorc顯存分配規(guī)律內(nèi)容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

最新評論