【深度学习实战（13）】训练框架之加载预训练权重

m0_51579041

已于 2024-05-01 09:30:30 修改

阅读量269

点赞数 1

分类专栏：深度学习实战文章标签：深度学习人工智能

于 2024-04-20 13:11:12 首次发布

本文链接：https://blog.csdn.net/m0_51579041/article/details/137996731

版权

深度学习实战专栏收录该内容

39 篇文章 2 订阅

订阅专栏

一、预训练权重加载流程

预训练权值加载流程：pretrained_dict -> yolo_dict(backbone) -> temp_dict(与model_dict匹配上的) -> model_dict(load_state_dict加载)

二、代码

# ------------------------------------------------------#
#   预训练权值加载流程：pretrained_dict -> yolo_dict(backbone)
#   -> temp_dict(与model_dict匹配上的) -> model_dict(load_state_dict加载)
#   -> print result
# ------------------------------------------------------#
if args.pretrained_weights != '':
    # ------------------------------------------------------#
    #   预训练权值文件
    # ------------------------------------------------------#
    if local_rank == 0:
        print('Load weights {}.'.format(args.pretrained_weights))

    # ----------------------#
    #   模型的key
    # ----------------------#
    model_dict = model.state_dict()
    # ----------------------#
    #   预训练权重的Key
    # ----------------------#
    pretrained_dict = torch.load(args.pretrained_weights, map_location=device)

    # ----------------------#
    #   通过yolo_dict获取预训练权重的backbone
    # ----------------------#
    yolo_dict = {}
    for k, v in pretrained_dict.items():
        if "backbone.backbone" in k:
            k = k.replace("backbone.backbone", "backbone")
        yolo_dict[k] = v

    # ----------------------#
    #   temp_dict用来更新model_dict
    #   load_key, no_load_key,用来记录哪些是加载了的
    # ----------------------#
    load_key, no_load_key, temp_dict = [], [], {}
    for k, v in yolo_dict.items():
        if k in model_dict.keys() and np.shape(model_dict[k]) == np.shape(v):
            temp_dict[k] = v
            load_key.append(k)
        else:
            no_load_key.append(k)
    # ----------------------#
    #   temp_dict用来更新model_dict
    # ----------------------#
    model_dict.update(temp_dict)
    model.load_state_dict(model_dict)
    # ------------------------------------------------------#
    #   打印没有匹配上的Key
    # ------------------------------------------------------#
    if local_rank == 0:
        print("\nSuccessful Load Key:", str(load_key)[:500], "……\nSuccessful Load Key Num:", len(load_key))
        print("\nFail To Load Key:", str(no_load_key)[:500], "……\nFail To Load Key num:", len(no_load_key))
        print("\n\033[1;33;44m温馨提示，head部分没有载入是正常现象，Backbone部分没有载入是错误的。\033[0m")