PyTorch加载预训练模型

Pytorch调用预训练模型：

1

Pytorch保存和加载整个模型：

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

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

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

2

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

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

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

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

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

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

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

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

	model = torchvision.models.densenet169()

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

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

若网络结构层严格对等：

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

网络结构层不对等：

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

参考：https://blog.csdn.net/aaon22357/article/details/82696938