修改权重使用预训练模型权重

最新推荐文章于 2023-10-31 15:02:05 发布

lsh呵呵

最新推荐文章于 2023-10-31 15:02:05 发布

阅读量980

点赞数

分类专栏： python与人工睿智机器学习入门与放弃文章标签：深度学习计算机视觉

≧◔◡◔≦ 如觉得有所帮助，请点个✌赞或关注✌。转载请注明出处。 ≧◔◡◔≦

本文链接：https://blog.csdn.net/nima1994/article/details/122132199

版权

python与人工睿智同时被 2 个专栏收录

146 篇文章 21 订阅

订阅专栏

机器学习入门与放弃

84 篇文章 18 订阅

订阅专栏

抛弃最后的输出层并非最佳方案。可以修改输出层的权重，以 mmdetection 使用预训练模型为例。

import torch

def faster_rcnn(num_classes):
    model_coco = torch.load("e:/14_model.pth")
    # weight
    for i in list(model_coco["state_dict"]):
        if 'bbox_head.fc_cls' in i:

            model_coco["state_dict"][i] = model_coco["state_dict"][i][:num_classes]
        elif 'bbox_head.fc_reg' in i:
            model_coco["state_dict"][i] = model_coco["state_dict"][i][:(num_classes * 4)]
    # save new model
    torch.save(model_coco, "../checkpoints/faster_rcnn_r50_fpn_1x_%d.pth" % num_classes)

def cascade_rcnn(num_classes):
    model_coco = torch.load("../checkpoints/cascade_rcnn_dconv_c3-c5_r101_fpn_1x_20190125-aaa877cc.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, "../checkpoints/cascade_rcnn_dconv_c3-c5_r101_fpn_1x_%d.pth" % num_classes)
    
# cascade_rcnn(num_classes = 11)
faster_rcnn(4)