pytorch：如何修改加载了预训练权重的模型的输入或输出--(修改torch.nn.DataParallel权重为cpu加载)

简介

如果训练模型使用的算法是在GPU上使用torch.nn.DataParallel加载多个GPU进行训练，那么是不可以直接在cpu上进行直接推理，原因是权重文件中的节点名称中均增加了一个module的参数文件。为了能再cpu上进行加载，除了需要注意使用跨设备加载时使用，map_location="cpu"之外，还需要对权重文件进行修改。

代码

import torch

def change_feature(check_point):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")   # 由于本文中是使用cpu，因此使用torch.load中将设备加载到cpu中，实际上可以直接使用torch.load进行加载，默认是cpu设备。
    check_point = torch.load(check_point, map_location=device)

    import collections
    dicts = collections.OrderedDict()

    for k, value in check_point.items():
        print("names:{}".format(k))                    # 打印结构
        print("shape:{}".format(value.size()))

        if "module" in k:            # 去除命名中的module
            k = k.split(".")[1:]
            k = ".".join(k)
            print k
        dicts[k] = value

    torch.save(dicts, "/home/zhaokaiyue/PycharmProjects/deepglobe/weights/log02_dink34.th")


if __name__ == "__main__":
    model_path = "/home/zhaokaiyue/PycharmProjects/deepglobe/weights/log01_dink34.th"
    change_feature(model_path)

原权重的结构参数

修改后的权重

可以看出，经过修改后的权重文件已经去掉节点名称中的module，此时进行测试发现，进行推理时进行正常。