最近遇到一個問題,發下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#纔是你的模型