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)