mmdetection 修改預訓練模型權重類別數

將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()

各個模型下載路徑

點擊這裏下載各個網絡模型

列如:

發佈了113 篇原創文章 · 獲贊 107 · 訪問量 23萬+
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章