pytorch篇---现有网络模型的使用和修改

以vgg16网络为例

1. 加载未预训练的网络

import torchvision
vgg16_false = torchvision.models.vgg16(pretrained=False)
print(vgg16_false)

2. 加载预预训练的网络

import torchvision
vgg16_true = torchvision.models.vgg16(pretrained=True)
print(vgg16_true)

网络结构查看

由此可见vgg16最后是全连接层 4096， 输出的是1000个类别

3. 加载数据（CIFAR10是10分类）

import torchvision
train_data = torchvision.datasets.CIFAR10('../data', train=True, transform=torchvision.transforms.ToTensor(),
                                          download=True)

4. 修改网络最后层的输出

由于vgg16是1000分类，要判别10分类，需要将网络的最后输出结果变成10分类。有以下两种办法：

办法一：在网络层的后面增加一个线程输出层

from torch import nn
vgg16_true.add_module('add_linear', nn.Linear(1000, 10))
print(vgg16_true)

控制台输出结果如下：
add_linear 加在classfier的外面，那如何加到里面呢？请看下一步

from torch import nn
vgg16_true.classifier.add_module('add_linear', nn.Linear(1000, 10))
print(vgg16_true)

办法二：直接将网络的最后一步4096 变成10分类

from torch import nn
vgg16_true.classifier[6] = nn.Linear(4096, 10)
print(vgg16_true)

5. 预测结果

for data in train_data:
    imgs, targets = data
    imgs = imgs.unsqueeze(0)
    output = vgg16_true(imgs)
    print(imgs.shape)
    print(targets)