pytorch 多GPU训练的模型在单卡或CPU上加载

在多GPU上并行训练的模型如何在单GPU或CPU上加载。

多卡训练的模型保存下来的参数名称是module.xxx.weight,而单卡的参数名称是xxx.weight，这时就会报错，找不到相应的字典的错误。

办法也很简单，只需要在state_dict上进行修改，将module.xxx.weight改为xxx.weight，再load就可以了。

代码如下：

import torch from collections import OrderedDict state_dict = torch.load(net_path) new_state_dict = OrderedDict() for k, v in state_dict.items(): # k为module.xxx.weight, v为权重     name = k[7:] # 截取`module.`后面的xxx.weight     new_state_dict[name] = v # load params net = XXXnet() net.load_state_dict(new_state_dict)