弄清Pytorch顯存的分配機(jī)制
對(duì)于顯存不充足的煉丹研究者來說，弄清楚Pytorch顯存的分配機(jī)制是很有必要的。下面直接通過實(shí)驗(yàn)來推出Pytorch顯存的分配過程。
實(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，這是合理的。另外，后面通過backward()，計(jì)算出x.grad，占存與x一致，所以最終一共占有顯存9GB，這也是合理的。但是，輸出顯示了顯存的峰值為12GB，這多出的3GB是怎么來的呢？首先畫出計(jì)算圖：
下面通過列表的形式來模擬Pytorch在運(yùn)算時(shí)分配顯存的過程：
如上所示，由于需要保存反向傳播以前所有前向傳播的中間變量，所以有了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運(yùn)算符，具有一定的參考價(jià)值，感興趣的小伙伴們可以參考一下，希望能夠給你帶來幫助2021-12-12
在本篇文章里，小編給大家分享的是一篇關(guān)于python爬蟲是否容易學(xué)的相關(guān)知識(shí)點(diǎn)內(nèi)容，有興趣的朋友們可以閱讀下。2020-06-06
Python復(fù)制Word內(nèi)容并使用格式設(shè)字體與大小實(shí)例代碼
這篇文章主要介紹了Python復(fù)制Word內(nèi)容并使用格式設(shè)字體與大小實(shí)例代碼，小編覺得還是挺不錯(cuò)的，具有一定借鑒價(jià)值，需要的朋友可以參考下2018-01-01
Python實(shí)現(xiàn)新型冠狀病毒傳播模型及預(yù)測(cè)代碼實(shí)例
在本篇文章里小編給大家整理的是關(guān)于Python實(shí)現(xiàn)新型冠狀病毒傳播模型及預(yù)測(cè)代碼內(nèi)容，有興趣的朋友們可以學(xué)習(xí)下。2020-02-02
Python創(chuàng)建一個(gè)元素都為0的列表實(shí)例
今天小編就為大家分享一篇Python創(chuàng)建一個(gè)元素都為0的列表實(shí)例，具有很好的參考價(jià)值，希望對(duì)大家有所幫助。一起跟隨小編過來看看吧2019-11-11
python讀取和保存為excel、csv、txt文件及對(duì)DataFrame文件的基本操作指南
最近在做一個(gè)項(xiàng)目,必須把結(jié)果保存到excel文件中,下面這篇文章主要給大家介紹了關(guān)于python讀取和保存為excel、csv、txt文件及對(duì)DataFrame文件的基本操作指南的相關(guān)資料,文中通過實(shí)例代碼介紹的非常詳細(xì),需要的朋友可以參考下2022-08-08
這篇文章主要介紹了如何利用Python編寫一個(gè)藏頭詩在線生成器，文中的示例代碼講解詳細(xì)，感興趣的同學(xué)可以跟隨小編一起動(dòng)手嘗試一下2022-04-04
Python json格式化打印實(shí)現(xiàn)過程解析
這篇文章主要介紹了Python json格式化打印實(shí)現(xiàn)過程解析,文中通過示例代碼介紹的非常詳細(xì)，對(duì)大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價(jià)值,需要的朋友可以參考下2020-07-07
python3使用urllib模塊制作網(wǎng)絡(luò)爬蟲
本文給大家介紹的是利用urllib模塊通過指定的URL抓取網(wǎng)頁內(nèi)容 所謂網(wǎng)頁抓取,就是把URL地址中指定的網(wǎng)絡(luò)資源從網(wǎng)絡(luò)流中讀取出來,保存到本地，有需要的小伙伴可以參考下2016-04-04
這篇文章主要介紹了Anaconda超詳細(xì)保姆級(jí)安裝配置教程,本篇教程對(duì)于新手很友好強(qiáng)烈推薦,看完這篇文章相信你對(duì)如何安裝配置Anaconda一定了然于胸,我們一起來看看吧2023-03-03