多GPU训练模型的调用------pytorch

上一讲讲到用pytorch进行单机多GPU进行模型的训练，但是此时的模型调用也将与单GPU有较大的区别，因为DataParallel虽然也是一个Pytorch的nn.Module，但是实际上有一个module的变量用来保存传入的实际模型，因此模型中含有module。在模型调用中，使用不同的保存方法将对应使用不同的调用方法。

1）如果如上讲中采用torch.save(model.state_dict(), file_name)的方式（注释部分），直接调用模型会出现错误：RuntimeError: Error(s) in loading state_dict for MymModle："module.outconv2.bias”等。

此时在调用模型进行转换时，即pth转pt时，需要去掉module的影响，具体如下所示：

    import torch
    from collections import OrderedDict
    from MyModel import MyModel

    device = torch.device('cpu') #使用cpu进行推理
    model = MyModel(1, 1)

    state_dict = torch.load('./../model/model.pth', map_location='cpu') #将模型转为CPU
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        namekey = k[7:]    # 去掉module的影响
        new_state_dict[namekey] = v
    model.load_state_dict(new_state_dict)

    #model = model.to(device)
    model.eval()  # 把模型转为test模式
    example = torch.rand(1,1,480,640)
    traced_net = torch.jit.trace(model, example)  # 转换模型
    traced_net.save("./../model/model.pt")    # 保存模型

2）如果如上讲中采用torch.save(model.module.state_dict(), file_name) 的方式（未注释部分），此时在调用模型进行转换时，即pth转pt时，则不需要去掉module的影响，具体如下所示：

    import torch
    from collections import OrderedDict
    from MyModel import MyModel

    device = torch.device('cpu') #使用cpu进行推理
    model = MyModel(1, 1)
    
    model.load_state_dict(torch.load('./../model/model_1013.pth', map_location='cpu'))  # 加载模型
    #model = model.to(device)
    model.eval()  # 把模型转为test模式
    example = torch.rand(1,1,480,640)
    traced_net = torch.jit.trace(model, example)  # 转换模型
    traced_net.save("./../model/model.pt")    # 保存模型

在保存模型过程中一定要注意保存的方式，这是关键。