問題描述

有時在載入已訓練好的模型時，會出現 out of memory 的錯誤提示，但仔細檢測使用的GPU卡並沒有再用且記憶體也沒有超出。

經查閱發現原來是訓練模型時使用的GPU卡和載入時使用的GPU卡不一樣導致的。個人感覺，因為pytorch的模型中是會記錄有GPU資訊的，所以有時使用不同的GPU載入時會報錯。

解決方法

gpu之間的相互轉換。即，將訓練時的gpu卡轉換為載入時的gpu卡。

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

當然，你也可以直接將載入模型時使用的gpu卡改為和訓練時的一樣。但在多人使用一個伺服器時，你想用的gpu卡已被使用，就需按上面方法轉換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',loc: storage)

gpu 1 --> gpu 0

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

補充知識：pytorch model.cuda()花費時間很長

如果pytorch在進行model.cuda(）操作需要花費的時間很長，長到你懷疑GPU的速度了，那就是不正常的。

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

以上這篇pytorch模型載入之gpu和cpu互轉操作就是小編分享給大家的全部內容了，希望能給大家一個參考，也希望大家多多支援我們。