Pytorch如何加载部分权重

if os.path.exists(self.load_params):
	pretext_model = torch.load(self.load_params)

w = pretext_model["fc2.weight"]
b = pretext_model["fc2.bias"]
pretext_model["fc2.weight"] = torch.cat((w,w,w,w[:2]),dim=0)
pretext_model["fc2.bias"] = torch.cat((b,b,b,b[:2]),dim=0)

self.net.load_state_dict(pretext_model)

if os.path.exists(self.load_params):
	pretext_model = torch.load(self.load_params) #加载的参数
	model_dict = net.state_dict()  #模型参数
	print(model_dict)
	print(pretext_model)

#如果模型有k层，就加载
state_dict = {k: v for k, v in pretext_model.items() if k in model_dict.keys()}
model_dict.update(state_dict)
net.load_state_dict(model_dict)

net = models.vgg19(pretrained=True) #下载与训练参数
print(net)  #查看网络结构
net.classifier[6] = torch.nn.Linear(4096,10) #将输出层修改为10分类

num_fc_ftr = net.classifier[6]
net.fc = nn.Linear(num_fc_ftr, 128)
net.out = nn.Linear(128, 10)