弄清Pytorch顯存的分配機(jī)制
對(duì)于顯存不充足的煉丹研究者來(lái)說(shuō),弄清楚Pytorch顯存的分配機(jī)制是很有必要的。下面直接通過(guò)實(shí)驗(yàn)來(lái)推出Pytorch顯存的分配過(guò)程。
實(shí)驗(yàn)實(shí)驗(yàn)代碼如下:
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的顯存創(chuàng)建變量x,然后計(jì)算y,再用y進(jìn)行反向傳播??梢钥吹?,創(chuàng)建x后與計(jì)算y后分別占顯存3GB與6GB,這是合理的。另外,后面通過(guò)backward(),計(jì)算出x.grad,占存與x一致,所以最終一共占有顯存9GB,這也是合理的。但是,輸出顯示了顯存的峰值為12GB,這多出的3GB是怎么來(lái)的呢?首先畫出計(jì)算圖:
下面通過(guò)列表的形式來(lái)模擬Pytorch在運(yùn)算時(shí)分配顯存的過(guò)程:
如上所示,由于需要保存反向傳播以前所有前向傳播的中間變量,所以有了12GB的峰值占存。
我們可以不存儲(chǔ)計(jì)算圖中的非葉子結(jié)點(diǎn),達(dá)到節(jié)省顯存的目的,即可以把上面的代碼中的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:
以上就是弄清Pytorch顯存的分配機(jī)制的詳細(xì)內(nèi)容,更多關(guān)于Pytorch 顯存分配的資料請(qǐng)關(guān)注腳本之家其它相關(guān)文章!
相關(guān)文章
python訓(xùn)練數(shù)據(jù)時(shí)打亂訓(xùn)練數(shù)據(jù)與標(biāo)簽的兩種方法小結(jié)
今天小編就為大家分享一篇python訓(xùn)練數(shù)據(jù)時(shí)打亂訓(xùn)練數(shù)據(jù)與標(biāo)簽的兩種方法小結(jié),具有很好的參考價(jià)值,希望對(duì)大家有所幫助。一起跟隨小編過(guò)來(lái)看看吧2018-11-11對(duì)tensorflow中tf.nn.conv1d和layers.conv1d的區(qū)別詳解
今天小編就為大家分享一篇對(duì)tensorflow中tf.nn.conv1d和layers.conv1d的區(qū)別詳解,具有很好的參考價(jià)值,希望對(duì)大家有所幫助。一起跟隨小編過(guò)來(lái)看看吧2020-02-02Python浮點(diǎn)型(float)運(yùn)算結(jié)果不正確的解決方案
這篇文章主要介紹了Python浮點(diǎn)型(float)運(yùn)算結(jié)果不正確的解決方案,幫助大家更好的利用python進(jìn)行運(yùn)算處理,感興趣的朋友可以了解下2020-09-09pytest用例執(zhí)行順序和跳過(guò)執(zhí)行詳解
本文主要介紹了pytest用例執(zhí)行順序和跳過(guò)執(zhí)行詳解,文中通過(guò)示例代碼介紹的非常詳細(xì),對(duì)大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價(jià)值,需要的朋友們下面隨著小編來(lái)一起學(xué)習(xí)學(xué)習(xí)吧2023-02-02Python中利用all()來(lái)優(yōu)化減少判斷的實(shí)例分析
在本篇文章里小編給大家整理的是一篇關(guān)于Python中利用all()來(lái)優(yōu)化減少判斷的實(shí)例分析內(nèi)容,有需要的朋友們可以學(xué)習(xí)下。2021-06-06Django中prefetch_related()函數(shù)優(yōu)化實(shí)戰(zhàn)指南
我們可以利用Django框架中select_related和prefetch_related函數(shù)對(duì)數(shù)據(jù)庫(kù)查詢優(yōu)化,這篇文章主要給大家介紹了關(guān)于Django中prefetch_related()函數(shù)優(yōu)化的相關(guān)資料,需要的朋友可以參考下2022-11-11Python+redis通過(guò)限流保護(hù)高并發(fā)系統(tǒng)
這篇文章主要介紹了Python+redis通過(guò)限流保護(hù)高并發(fā)系統(tǒng),文中通過(guò)示例代碼介紹的非常詳細(xì),對(duì)大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價(jià)值,需要的朋友可以參考下2020-04-04python爬蟲指南之xpath實(shí)例解析(附實(shí)戰(zhàn))
在進(jìn)行網(wǎng)頁(yè)抓取的時(shí)候,分析定位html節(jié)點(diǎn)是獲取抓取信息的關(guān)鍵,目前我用的是lxml模塊,下面這篇文章主要給大家介紹了關(guān)于python爬蟲指南之xpath實(shí)例解析的相關(guān)資料,需要的朋友可以參考下2022-01-01