科研小白进阶之路——torch实现模型保存和加载的几种方式

模型的保存：

方式一：保存整个模型

model = resnet()
#-------省略训练过程------------
torch.save(model,path)

方式二：只保存模型参数

model = resnet()
#-------省略训练过程---------------
torch.save(model.state_dict,path)

模型的加载：
情况一：训练已经结束，测试时加载训练好的模型
1、在训练的py文件下面直接对模型进行测试：

model = resnet()     # 定义模型
#----训练过程省略-----
torch.save(model,path) #保存模型
#----训练好的模型进行测试----
model.eval
#---加载测试集等步骤

2、单独写一个infer.py
①加载整个模型和参数，对应模型的保存方式一，即当时就保存了整个模型，此时直接加载，然后去测试；

model_path = './model_000009.pth'
model = torch.load(model_path)   # 加载了整个模型

②只加载模型参数，对应模型保存的方式二，当时只保存了模型参数，此时需要先初始化模型，然后让模型加载训练好的模型的参数；

model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))

情况二：训练一半停了，想继续训练
只加载模型参数，网络结构从代码中创建

net = resnet()  #代码中创建网络结构
path = './pth'
saved_state_dict = torch.load(path) #加载模型参数
net.load_state_dict(saved_state_dict) #应用到网络中

情况三：加载预训练模型
1、使用torchvision自带的预训练模型

#----------resnet系列------------
import torchvision.models as models
model = models.ResNet(pretrained=True)
model = models.resnet18(pretrained=True)
model = models.resnet34(pretrained=True)
model = models.resnet50(pretrained=True)
#---------VGG系列-------------
model = models.VGG(pretrained=True)
model = models.vgg11(pretrained=True)
model = models.vgg16(pretrained=True)
model = models.vgg16_bn(pretrained=True)
#---------alexnet-----------
Alexnet = models.alexnet()
#--------squeezenet--------
squeezenet = models.squeezenet1_0()

2、如果只需要加载预训练模型的部分参数：
根据名称进行加载

import torchvision.models as models 
model = models.resnet18(pretrained=True)
pretrained_dict = model.state_dict() #加载torchvision中的预训练模型和参数后通过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.update(pretrained_dict) #更新model_dict
model.load_state_dict(model_dict)

简单参数修改

import torchvision.models as models
model = models.resnet50(pretrained=True)
fc_features = model.fc.in_features
model.fc = nn.Linear(fc_features, class_num)

增删卷积层

import torchvision.models as models
import torch
import torch.nn as nn
import math
import torch.utils.model_zoo as model_zoo

class CNN(nn.module):
	def _init_(self,block,layers,num_classes):
		self.inplanes = 64
		super(Resnet,self)._init_()
		self.conv1 = nn.Conv2d(3,64,kernel_size=7,stride=2, padding=3,bias=False)
		self.bn1 = nn.BatchNorm2d(64)
		self.relu = nn.Relu(inplace=True)
		self.maxpool = self.MaxPool2d(kernelsize=3,stride=2,padding=1)
		self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.avgpool = nn.AvgPool2d(7, stride=1)
        #新增一个反卷积层
        self.convtranspose1 = nn.ConvTranspose2d(2048, 2048, kernel_size=3, stride=1, padding=1, output_padding=0, groups=1, bias=False, dilation=1)
        #新增一个最大池化层
        self.maxpool2 = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
        #去掉原来的fc层，新增一个fclass层
        self.fclass = nn.Linear(2048, num_classes)
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
	def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )
 
        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))
 
        return nn.Sequential(*layers)
	def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
 
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
 
        x = self.avgpool(x)
        #新加层的forward
        x = x.view(x.size(0), -1)
        x = self.convtranspose1(x)
        x = self.maxpool2(x)
        x = x.view(x.size(0), -1)
        x = self.fclass(x)
 
        return x

#加载model
resnet50 = models.resnet50(pretrained=True)
cnn = CNN(Bottleneck, [3, 4, 6, 3])
#读取参数
pretrained_dict = resnet50.state_dict()
model_dict = cnn.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
cnn.load_state_dict(model_dict)

ResNet预训练模型下载地址

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',
}

参考博客
https://www.cnblogs.com/wangguchangqing/p/11058525.html.
https://blog.csdn.net/whut_ldz/article/details/78845947.