有时无法直接利用官方的预训练模型,需要自己手动将指定模型进行预训练。主要有以下语句作为参考:
model = models.resnet34(pretrained=False)
pretrained_dict = torch.load('./pretrain/resnet34-333f7ec4.pth')
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model.state_dict()}
model.load_state_dict(pretrained_dict)
model.fc = torch.nn.Linear(512, 5) # 512为原始fc的数目,5是自己任务的分类数
2、更新权重
model_dict = model.state_dict()
# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, vin pretrained_dict.items() if k in model_dict}
model_dict.update(pretrained_dict)



