pytorch 如何加载部分预训练模型

pretrained_dict =...

model_dict = model.state_dict()


# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, vin pretrained_dict.items() if k inmodel_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)

你可能感兴趣的:(pytorch)