2019-04-02 Pytorch:加载预训练模型并保存

 参考链接:

1.用pytorch搭建AlexNet(微调预训练模型及手动搭建) - sjtuxx_lee的博客 - CSDN博客

2.pytorch学习笔记之加载预训练模型 - spectre - CSDN博客


第一篇文章实现了直接加载Alexnet预训练模型,并根据自己的需要微调(将最后一层全连接层输出由1000改为10)。运行没有问题,下面是我自己的理解和备注。原代码可以在参考链接2中直接copy。

import torch.nn as nn                                                 #nn设置网络结构详细参数

from torchvision import models                                  #torchvision包,它包括3个子包,分别是: torchvison.datasets ,torchvision.models ,torchvision.transforms ,分别是预定义好的数据集(比如MNIST、CIFAR10等)、预定义好的经典网络结构(比如AlexNet、VGG、ResNet等)和预定义好的数据增强方法(比如Resize、ToTensor等)。

class BuildAlexNet(nn.Module):

    def __init__(self, model_type, n_output):

        super(BuildAlexNet, self).__init__()

        self.model_type = model_type

        if model_type == 'pre':                                         #定义两种model类型,一个直接从alexnet中继承这个参数和结构,定义名称为‘pre’。 另一个是自己设定的网络结构,定义为'new'

            model = models.alexnet(pretrained=True)       #加载alexnet模型,pretrained为真,则加载网络结构和预训练参数。否则,只加载网络结构[2]

            self.features = model.features                          #因为只要求更改最后的分类数,所以feature类直接从预训练网络中继承classifier   类除了要更改的分类层,其他的也从原网络中定义好

            fc1 = nn.Linear(9216, 4096)                                #fc1和fc2继承原网络的classifier参数

            fc1.bias = model.classifier[1].bias

            fc1.weight = model.classifier[1].weight

            fc2 = nn.Linear(4096, 4096)

            fc2.bias = model.classifier[4].bias

            fc2.weight = model.classifier[4].weight

            self.classifier = nn.Sequential(                          #定义新的classifier层,前两层保持不变,底端分类层分类数用n_output代替

                    nn.Dropout(),

                    fc1,

                    nn.ReLU(inplace=True),

                    nn.Dropout(),

                    fc2,

                    nn.ReLU(inplace=True),

                    nn.Linear(4096, n_output)) 

            #或者直接修改为

#            model.classifier[6]==nn.Linear(4096,n_output)

#            self.classifier = model.classifier

        if model_type == 'new':                                         #这是自己定义的网络模型(feature,classifier)

            self.features = nn.Sequential(

                    nn.Conv2d(3, 64, 11, 4, 2),

                    nn.ReLU(inplace = True),

                    nn.MaxPool2d(3, 2, 0),

                    nn.Conv2d(64, 192, 5, 1, 2),

                    nn.ReLU(inplace=True),

                    nn.MaxPool2d(3, 2, 0),

                    nn.Conv2d(192, 384, 3, 1, 1),

                    nn.ReLU(inplace = True),

                    nn.Conv2d(384, 256, 3, 1, 1),

                    nn.ReLU(inplace=True),

                    nn.MaxPool2d(3, 2, 0))

            self.classifier = nn.Sequential(

                    nn.Dropout(),

                    nn.Linear(9216, 4096),

                    nn.ReLU(inplace=True),

                    nn.Dropout(),

                    nn.Linear(4096, 4096),

                    nn.ReLU(inplace=True),

                    nn.Linear(4096, n_output))

    def forward(self, x):

        x = self.features(x)

        x = x.view(x.size(0), -1)

        out  = self.classifier(x)

        return out


import numpy as np                                                       #从这里开始就是验证函数

from torch.autograd import Variable

import torch

if __name__ == '__main__':

    model_type = 'pre'

    n_output = 10

    alexnet = BuildAlexNet(model_type, n_output)     #调用函数buildAlexnet,网络选项是预训练模型,输出是10,也就是分十类

    print(alexnet)

    x = np.random.rand(1,3,224,224)              #随机创建一个224*224,通道为3的数组,模拟三通道的图片

    x = x.astype(np.float32)

    x_ts = torch.from_numpy(x)                     #转换成torch能用的张量形式

    x_in = Variable(x_ts)                                  #『PyTorch』第五弹_深入理解autograd_上:Variable属性方法 - 叠加态的猫 - 博客园

    y = alexnet(x_in)                                         #输出结果


打印出来的网络结构如下:


2019-04-02 Pytorch:加载预训练模型并保存_第1张图片

输出分类结果y如下:

最后的分类结果是10类,check~

待解决:1.代码倒数第二行的torch.autograd.variable没有理解清楚,链接已贴.上. ,等您脑袋清楚了一定要看哈。

               2.继承,super(),代码第五行

你可能感兴趣的:(2019-04-02 Pytorch:加载预训练模型并保存)