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

ゞ浴缸里的玫瑰 2023-02-20 06:16 213阅读 0赞

## 简介 ##

如果训练模型使用的算法是在`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)

**原权重的结构参数**  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5ODUyNjc2_size_16_color_FFFFFF_t_70]  
**修改后的权重**  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5ODUyNjc2_size_16_color_FFFFFF_t_70 1]  
可以看出，经过修改后的权重文件已经去掉节点名称中的`module`，此时进行测试发现，进行推理时进行正常。

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5ODUyNjc2_size_16_color_FFFFFF_t_70]: https://img-blog.csdnimg.cn/20200623173753360.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5ODUyNjc2,size_16,color_FFFFFF,t_70
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5ODUyNjc2_size_16_color_FFFFFF_t_70 1]: https://img-blog.csdnimg.cn/20200623173831874.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5ODUyNjc2,size_16,color_FFFFFF,t_70