腳本之家服務(wù)器常用軟件

快捷導(dǎo)航

如何解決pytorch訓(xùn)練過(guò)程中CPU內(nèi)存溢出問(wèn)題

更新時(shí)間：2023年09月08日 09:15:59 作者：yyyyyyhx

這篇文章主要介紹了如何解決pytorch訓(xùn)練過(guò)程中CPU內(nèi)存溢出問(wèn)題,具有很好的參考價(jià)值,希望對(duì)大家有所幫助,如有錯(cuò)誤或未考慮完全的地方,望不吝賜教

pytorch訓(xùn)練過(guò)程中CPU內(nèi)存溢出問(wèn)題

期待了一晚上的結(jié)果，早上起來(lái)竟然發(fā)現(xiàn)CPU內(nèi)存溢出導(dǎo)致程序奔潰了，今天郁悶了一天。。。

經(jīng)查詢，內(nèi)存溢出的常見(jiàn)原因有：

找了好多原因，感覺(jué)上面的比較貼近，但是改了一堆，內(nèi)存還是蹭蹭往上漲。

后來(lái)調(diào)用了memory_profiler這個(gè)包，發(fā)現(xiàn)了

我的程序中上面這部分每輪都加了70多m，可能是這個(gè)循環(huán)導(dǎo)致的（另外排查了好多程序都不說(shuō)了。。。）

然后我把這個(gè)循環(huán)改成了下面這樣

觀察每組數(shù)據(jù)的memory占用，發(fā)現(xiàn)了很神奇的現(xiàn)象：

第一組數(shù)據(jù)沒(méi)怎么占內(nèi)存，后面重復(fù)調(diào)用同一個(gè)數(shù)據(jù)的時(shí)候內(nèi)存就開(kāi)始暴漲，到現(xiàn)在也沒(méi)整明白為什么會(huì)這樣。。

后來(lái)調(diào)用了

from einops import rearrange

把數(shù)組維度修改了一下，再一起送進(jìn)網(wǎng)絡(luò)

可算解決了

pytorch顯存越來(lái)越多的一個(gè)原因

optimizer.zero_grad()
loss.backward()
optimizer.step()
train_loss += loss

參考了別人的代碼發(fā)現(xiàn)那句loss一般是這樣寫

loss_sum += loss.data[0]

這是因?yàn)檩敵龅膌oss的數(shù)據(jù)類型是Variable。而PyTorch的動(dòng)態(tài)圖機(jī)制就是通過(guò)Variable來(lái)構(gòu)建圖。

主要是使用Variable計(jì)算的時(shí)候，會(huì)記錄下新產(chǎn)生的Variable的運(yùn)算符號(hào)，在反向傳播求導(dǎo)的時(shí)候進(jìn)行使用。

如果這里直接將loss加起來(lái)，系統(tǒng)會(huì)認(rèn)為這里也是計(jì)算圖的一部分，也就是說(shuō)網(wǎng)絡(luò)會(huì)一直延伸變大那么消耗的顯存也就越來(lái)越大。

用Tensor計(jì)算要寫成：

train_loss += loss.item()

correct_total += torch.eq(predict, label_batch).sum().item()
train_loss += loss.item()

當(dāng)需要將模型中變量提取出來(lái)參與計(jì)算時(shí)，需要使用** .item()**

以上為個(gè)人經(jīng)驗(yàn)，希望能給大家一個(gè)參考，也希望大家多多支持腳本之家。

您可能感興趣的文章: