pytorch 模型加载部分权重与部分模型加载权重

1. 模型加载部分权重

model = ...
model_dict = model.state_dict()
pretrained_dict = torch.load(load_name)
# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, v in model_dict.items() if k in pretrained_dict}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
# 3. load the new state dict
model.load_state_dict(model_dict)

2. 部分模型加载权重

model = ...
model_dict = model.state_dict()
pretrained_dict = torch.load(load_name)
# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
# 3. load the new state dict
model.load_state_dict(model_dict)