pytorch 自己选择模块加载权重

1. 打印一个权重模型的节点名称

pretrained_dict = torch.load(r"logs_c_3/LPRNet__iteration_70000.pth",
                             map_location=device)
for k, v in pretrained_dict.items():
    print(k)

2. 加载匹配上的权重

lprnet = build_lprnet(lpr_max_len=args.lpr_max_len, phase=args.phase_train, class_num=len(CHARS), class_num_color=3,dropout_rate=args.dropout_rate)
device = torch.device("cuda:0" if args.cuda else "cpu")
lprnet.to(device)

if pretrained_model:
	lprnet.load_state_dict(torch.load(pretrained_model),strict=False)  # strict为False时，节点名称不匹配就不加载

3. 自己选择权重加载

if pretrained_model:
	model_dict = lprnet.state_dict()
	pretrained_dict = torch.load(pretrained_model,map_location=device)
	pretrained_dict = {k: v for k, v in pretrained_dict.items() if
	                  k.split('.')[0] in ['backbone', 'character_layers', 'container']}
	# pretrained_dict = {k: v for k, v in pretrained_dict.items() if
	#                   k.split('.')[0] not in ['character_layers', 'container']}
	
	model_dict.update(pretrained_dict)
	lprnet.load_state_dict(model_dict)