pytorch GPU和CPU模型相互加载

前言

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

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

torch.load('model_dict.pkl')

情况二:模型是CPU,预加载的训练参数却是GPU,那么需要这样:

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

情况三:模型是GPU,预加载的训练参数却是CPU:

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

猜你喜欢

转载自blog.csdn.net/u012135425/article/details/85217542