問(wèn)題描述

有時(shí)在加載已訓(xùn)練好的模型時(shí)，會(huì)出現(xiàn) out of memory 的錯(cuò)誤提示，但仔細(xì)檢測(cè)使用的GPU卡并沒(méi)有再用且內(nèi)存也沒(méi)有超出。

經(jīng)查閱發(fā)現(xiàn)原來(lái)是訓(xùn)練模型時(shí)使用的GPU卡和加載時(shí)使用的GPU卡不一樣導(dǎo)致的。個(gè)人感覺(jué)，因?yàn)閜ytorch的模型中是會(huì)記錄有GPU信息的，所以有時(shí)使用不同的GPU加載時(shí)會(huì)報(bào)錯(cuò)。

解決方法

gpu之間的相互轉(zhuǎn)換。即，將訓(xùn)練時(shí)的gpu卡轉(zhuǎn)換為加載時(shí)的gpu卡。

torch.load('modelparameters.pth', map_location={'cuda:1':'cuda:0'}) # gpu 1 --> gpu 0

當(dāng)然，你也可以直接將加載模型時(shí)使用的gpu卡改為和訓(xùn)練時(shí)的一樣。但在多人使用一個(gè)服務(wù)器時(shí)，你想用的gpu卡已被使用，就需按上面方法轉(zhuǎn)換gpu。

拓展

cpu --> cpu 或 gpu --> gpu

checkpoint = torch.load('modelparameters.pth')
model.load_state_dict(checkpoint)

cpu --> gpu 0

torch.load('modelparameters.pth', map_location=lambda storage, loc: storage.cuda(0))

gpu --> cpu

torch.load('modelparameters.pth', map_location=lambda storage, loc: storage)

gpu 1 --> gpu 0

torch.load('modelparameters.pth', map_location={'cuda:1':'cuda:0'})

補(bǔ)充知識(shí)：pytorch model.cuda()花費(fèi)時(shí)間很長(zhǎng)

如果pytorch在進(jìn)行model.cuda(）操作需要花費(fèi)的時(shí)間很長(zhǎng)，長(zhǎng)到你懷疑GPU的速度了，那就是不正常的。

如果你用的pytorch版本是0.3.0，升級(jí)到0.3.1就好了！

以上這篇pytorch模型載入之gpu和cpu互轉(zhuǎn)操作就是小編分享給大家的全部?jī)?nèi)容了，希望能給大家一個(gè)參考，也希望大家多多支持腳本之家。

最新評(píng)論