一、预训练权重加载流程

预训练权值加载流程：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] = vload_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")