Pytorch调用预训练模型

1.Pytorch保存和加载整个模型：

  1. torch.save(model, 'model.pth')
  2. model = torch.load('model.pth')

2.Pytorch保存和加载预训练模型参数：

  1. torch.save(model.state_dict(), 'params.pth')
  2. model.load_state_dict(torch.load('params.pth'))

3.Pytorch调用网络提供的预训练模型：

torchvision.models这个包中包含alexnet、densenet、inception、resnet、 squeezenet、vgg等常用的网络结构，并且提供了预训练模型，可以通过简单调用来读取网络结构和预训练模型。

  1. import torchvision 
  2. model = torchvision.models.resnet50(pretrained=True)

这样就导入了resnet50的预训练模型了。

如果只需要网络结构，不需要用预训练模型的参数来初始化，那么就是:

  1. model = torchvision.models.resnet50(pretrained=False)

如果要导入densenet模型也是同样的道理，比如导入densenet169，且不需要是预训练的模型：

  1. model = torchvision.models.densenet169(pretrained=False)

由于pretrained参数默认是False，所以等价于：

  1. model = torchvision.models.densenet169()

4.下载预训练模型

接下来以导入resnet50为例介绍具体导入模型时候的源码。运行model = torchvision.models.resnet50(pretrained=True)的时候，是通过models包下的resnet.py脚本进行的，源码如下：

  1. import torch.nn as nn
  2. import math
  3. import torch.utils.model_zoo as model_zoo   # model_zoo是和导入预训练模型相关的包
  4. __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
  5.            'resnet152']
  6. # model_urls这个字典是预训练模型的下载地址
  7. model_urls = {    
  8.     'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
  9.     'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
  10.     'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
  11.     'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
  12.     'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
  13. }

1)若网络结构层严格对等：

  1. def resnet50(pretrained=False, **kwargs):
  2.     model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
  3.     if pretrained:
  4.         model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))

2)网络结构层不对等：

  1. def resnet50_cbam(pretrained=False, **kwargs):
  2.     model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
  3.     if pretrained:
  4.         pretrained_state_dict = model_zoo.load_url(model_urls['resnet50']) #load_url函数根据model_urls字典下载或导入相应的预训练模型
  5.         now_state_dict        = model.state_dict()   # 返回model模块的字典
  6.         now_state_dict.update(pretrained_state_dict)
  7.         model.load_state_dict(now_state_dict) #最后通过调用model的load_state_dict方法用预训练的模型参数来初始化你构建的网络结构，这个方法就是PyTorch中通用的用一个模型的参数初始化另一个模型的层的操作。load_state_dict方法还有一个重要的参数是strict，该参数默认是True，表示预训练模型的层和你的网络结构层严格对应相等（比如层名和维度）
  8.     return model

转载至：https://blog.csdn.net/vino_cherish/article/details/84110401