加载预训练模型——torch.load 与 load_state_dict
程序员文章站
2022-06-13 15:57:52
...
checkpoint = torch.load('./pretrained/Clean_resnet18-200-regular.pth')
# torch.load加载预训练模型参数
model.load_state_dict(checkpoint['state_dict'])
# load_state_dict是net的一个方法：把torch.load加载的参数加载到模型中
model.eval()
# 把模型换成测试模式。因为测试模型的时候不用dropout和BN了