pytorch加载模型

加载预训练模型

1 2	my_resnet = MyResNet(args, *kwargs) my_resnet.load_state_dict(torch.load("my_resnet.pth"))

1	torch.save(my_resnet.state_dict(), "my_resnet.pth")

1	my_resnet = torch.load("my_resnet.pth")

加载部分预训练模型

pretrained_dict = model_zoo.load_url(model_urls['resnet152'])
model_dict = model.state_dict()
# 将pretrained_dict里不属于model_dict的键剔除掉
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 更新现有的model_dict
model_dict.update(pretrained_dict)
# 加载我们真正需要的state_dict
model.load_state_dict(model_dict)