pytorch模型与参数的GPU,cpu加载转换

我们在使用pytorch的过程,经常会需要加载模型参数,不管是别人提供给我们的模型参数,还是我们自己训练的模型参数,那么加载模型参数就会碰到一些情况,即GPU模型和CPU模型,这两种模型是不能混为一谈的,下面分情况进行操作说明。

情况一:模型是GPU模型,预加载的训练参数也是GPU;模型是CPU模型,预加载的训练参数也是CPU,这种情况下我们都只用直接用下面的语句即可:

torch.load('model_dict.pkl')
情况二:GPU->CPU 模型是CPU,预加载的训练参数却是GPU,那么需要这样:

torch.load('model_dict.pkl', map_location=lambda storage, loc: storage)
情况三:CPU->GPU 模型是GPU,预加载的训练参数却是CPU:

torch.load('model_dic.pkl', map_location=lambda storage, loc: storage.cuda)

#CPU->GPU1   模型是GPU1,预加载的训练参数却是CPU:
torch.load('model_dic.pkl', map_location=lambda storage, loc: storage.cuda(1))
原文地址:https://www.cnblogs.com/Henry-ZHAO/p/14341662.html