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

一、预训练权重加载流程

预训练权值加载流程：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")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53