我们在使用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))