pytorch加载预训练模型

加载预训练模型

1/4 直接加载预训练模型

import torchvison.models as models
resnet50 = models.resnet50(pretrained=True)
# 若只需要网络结构，不需要预训练模型的参数来初始化，则设置pretrained=False，或者把resnet复制到自己的目录下，新建个model文件夹，可以参考下面的猫狗大战入门算法入门：https://github.com/JackwithWilshere/Kaggle-Dogs_vs_Cats_PyTorch

2/4 修改某一层

# resnet默认是ImageNet的1000分类
# 更改为二分类
resnet.fc = nn.Linear(2048, 2)

# resne第一层卷积的卷积核是7
# 更改第一层卷积的卷积核为5，需要计算featuremap的维度
resnet.conv1 == nn.Conv2d(3, 64, kernel_size=5, stride=2, padding=3, bias=False)

3/4 加载部分预训练模型

# 对于具体的任务，很难保证模型和公开的模型完全一样，但是预训练模型的参数确实有助于提高训练的准确率，为了结合二者的优点，需要加载部分预训练模型

#加载model，model是自己定义好的模型
resnet50 = models.resnet50(pretrained=True) 
model =Net(...) 
 
#读取参数 
pretrained_dict =resnet50.state_dict() 
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)  

4/4 加载自己的模型

# 即保存和恢复模型：训练好的模型保存，然后加载用于测试

# method 1 : （官方推荐）只保存和回复模型中的参数
# 使用这种方法，需要自己导入模型的结构信息
# step 1/2 : 保存
torch.save(model.state_dict(), PATH)
# example
torch.save(renet50.state_dict(), 'ckp/model.pth')
# setp 2/2 : 恢复
model = ModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
# example
resnet = resnet50(pretrained=True)
resnet.load_state_dict(troch.load'cpk/model.pth')

# method 2 : 保存模型的参数和结构信息
# step 1/2 : 保存
torch.save(model, PATH)
# setp 2/2 : 恢复
model = torch.load(PATH)