亚洲乱码中文字幕综合,中国熟女仑乱hd,亚洲精品乱拍国产一区二区三区,一本大道卡一卡二卡三乱码全集资源,又粗又黄又硬又爽的免费视频

弄清Pytorch顯存的分配機(jī)制

 更新時(shí)間:2020年12月10日 10:18:22   作者:頎周  
這篇文章主要介紹了Pytorch顯存的分配機(jī)制的相關(guān)資料,幫助大家更好的理解和使用Pytorch,感興趣的朋友可以了解下

  對(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)文章

最新評(píng)論