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

一、代码

import config
import torch
import numpy as np

device = 'cpu'
# ------------------------------------------------------#
#   创建你的模型
# ------------------------------------------------------#
model = YourModel()
if 0:
    weights_init(model)

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

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

    # ----------------------#
    #   temp_dict用来更新model_dict
    #   load_key, no_load_key,用来记录哪些是加载了的
    # ----------------------#
    load_key, no_load_key, temp_dict = [], [], {}
    for k, v in pretrained_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
    # ------------------------------------------------------#
    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")

运行结果