mmdetection 修改预训练模型权重类别数

ゝ一世哀愁。 2023-06-27 13:59 12阅读 0赞

**将coco预训练模型类别权重修改类训练数据集权重，代码如下:**  
如何修改可以看key里面值进行修改：

for key, value in model_coco["state_dict"].items():
            print(key)

# -*- coding: utf-8 -*-
    # @Time    : 20-1-7 上午10:28
    # @Author  : wusaifei
    # @FileName: Modify_category.py
    # @Software: PyCharm
    def main():
        #gen coco pretrained weight
        import torch
        num_classes = 11
        model_coco = torch.load("../checkpoints/cascade_rcnn_r50_fpn_1x_20190501-3b6211ab.pth") # weight
        for key, value in model_coco["state_dict"].items():
            print(key)
    ######################################################################################
        # faster_rcnn_r50_fpn_1x_20181010-3d1b3351.pth
        # model_coco["state_dict"]["bbox_head.fc_cls.weight"] = \
        # model_coco["state_dict"]["bbox_head.fc_cls.weight"][:num_classes, :]
        #
        #
        # model_coco["state_dict"]["bbox_head.fc_cls.bias"] = \
        # model_coco["state_dict"]["bbox_head.fc_cls.bias"][:num_classes]
    ###################################################################################### 
        # cascade_rcnn_r50_fpn_1x_20190501-3b6211ab.pth
        model_coco["state_dict"]["bbox_head.0.fc_cls.weight"] = \
        model_coco["state_dict"]["bbox_head.0.fc_cls.weight"][:num_classes, :]
    
        model_coco["state_dict"]["bbox_head.1.fc_cls.weight"] = \
        model_coco["state_dict"]["bbox_head.1.fc_cls.weight"][:num_classes, :]
    
        model_coco["state_dict"]["bbox_head.2.fc_cls.weight"] = \
        model_coco["state_dict"]["bbox_head.2.fc_cls.weight"][:num_classes, :]
    
        model_coco["state_dict"]["bbox_head.0.fc_cls.bias"] = \
        model_coco["state_dict"]["bbox_head.0.fc_cls.bias"][:num_classes]
    
        model_coco["state_dict"]["bbox_head.1.fc_cls.bias"] = \
        model_coco["state_dict"]["bbox_head.1.fc_cls.bias"][:num_classes]
    
        model_coco["state_dict"]["bbox_head.2.fc_cls.bias"] = \
        model_coco["state_dict"]["bbox_head.2.fc_cls.bias"][:num_classes]
        # save new model
        torch.save(model_coco, "cascade_rcnn_r50_fpn_1x_coco_pretrained_weights_classes_%d.pth" % num_classes)
    if __name__ == "__main__":
        main()

# 各个模型下载路径 #

### [点击这里下载各个网络模型][Link 1] ###

![列如：][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM4NDEwNDI4_size_16_color_FFFFFF_t_70]

[Link 1]: https://zhuanlan.zhihu.com/p/101216931
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM4NDEwNDI4_size_16_color_FFFFFF_t_70]: https://img-blog.csdnimg.cn/20200113094729580.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM4NDEwNDI4,size_16,color_FFFFFF,t_70