欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

PyTorch使用cpu调用gpu训练的模型

程序员文章站 2024-01-25 20:02:40
...

最近遇到一个问题,发下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#才是你的模型
相关标签: PyTorch 采坑