【pytorch】gpu訓練的模型無法在cpu下使用問題解決

最近遇到一個問題，發下gpu訓練的模型在cpu下不能跑，很難受。
結果發現原來是自己使用了DataParallel的問題，我保存的是整個模型。

model=DataParalle(model)
torch.save(‘xx.pkl’,model)

然後把模型copy到自己電腦上用cpu跑

model=torch.load(‘xx.pkl’)
#此處報錯了。。。。

原因是cpu環境不能直接導入gpu訓練的DataParallel模型

所以換個策略,現在gpu把模型轉化掉

model = DataParallel(model)
…
real_model = model.module#這個纔是你實際的模型，如果直接報錯model的話，其實是保存了DataParallel(model)這個，這樣會導致cpu環境下加載出錯
torch.save(real_model,‘xxx.pkl’) #這樣纔是正確的保存模型方式，這樣在cpu環境的模型纔不會出錯

其實還有一種方案，下面這種

model = torch.load(‘xxx.pkl’, map_location=lambda storage, loc: storage)
model = model.module#纔是你的模型

#上面那種貌似已經過時了,0.4.0的新方式

model = torch.load(‘xxx.pkl’, map_location=‘cpu’)
model = model.module#纔是你的模型