Pytorch官网上给我们提供了很多已训练好或未训练的网络模型,我们可以直接来进行使用
import torchvision.models as models
resnet18 = models.resnet18()
alexnet = models.alexnet()
vgg16 = models.vgg16()
squeezenet = models.squeezenet1_0()
densenet = models.densenet161()
inception = models.inception_v3()
googlenet = models.googlenet()
shufflenet = models.shufflenet_v2_x1_0()
mobilenet_v2 = models.mobilenet_v2()
mobilenet_v3_large = models.mobilenet_v3_large()
mobilenet_v3_small = models.mobilenet_v3_small()
resnext50_32x4d = models.resnext50_32x4d()
wide_resnet50_2 = models.wide_resnet50_2()
mnasnet = models.mnasnet1_0()
接下来我们使用VGG16model来进行详细讲解
import torchvision.models as models
vgg16=models.vgg16(pretrain=True)
print(vgg16)
当pretrain预训练设置为True的时候,我们将从网络中下载已训练好的模型,然后我们将其打印出来看一下它的结构。
我们可以看到,最后线性层的输出为1000,但是如果我们想用这个模型进行10分类该怎么办,就像运用到CIFAR10数据集上,我们只需要在最后一层的线性层后加一个层。
from torch import nn
vgg16.add_module('add_ten_classify_linear',nn.Linear(1000,10))
print(vgg16)
add_module函数需要两个参数,第一个是这个层的名字,我们自己定义就好,第二个是我们要添加的层,接着我们打印一下看看是否成功添加进去。
我们发现已经添加进去了,但是是在模型的最后,如果我们想添加到classifier里面该如何操作呢。
vgg16.classifier.add_module('add_ten_classify_linear',nn.Linear(1000,64))
这时我们可以看到该层已经添加到classifier里面了。
那如果我们想要修改原模型中的参数该如何操作呢。其实可以简单粗暴一点。
vgg16.classifier[6] = nn.Linear(in_features=4096,out_features=10,bias=True)
这样,我们直接将classifier中的第七层改为10输出的层。
还是简单做个记录,方便以后查阅。