pytorch入门17：现有模型的使用和修改_pytorch将符合的模块改为其他的

import ssl
"""
这一节主要是学会怎么在现有模型上进行修改模型结构
"""
import torchvision

# ssl._create_default_https_context = ssl._create_unverified_context  # 全局取消证书验证
# dataset = torchvision.datasets.ImageNet('../imagenet_data',split='train',download=True,transform=torchvision.transforms.ToTensor())
from torch import nn
from torch.utils.data import DataLoader

vgg16_fal = torchvision.models.vgg16(pretrained=False) #还需要训练的模型
vgg16_tr = torchvision.models.vgg16(pretrained=True)  #已经训练好的模型
#print('ok')
#print(vgg16_tr)  # vgg16是一个对1000种分类模型

# 现在想将vgg16用到CIFAR10上面，即改为10分类模型
dataset = torchvision.datasets.CIFAR10('./torchvision_dataset',train=False,download=False,transform=torchvision.transforms.ToTensor())
dataloader = DataLoader(dataset,batch_size=64,drop_last=True)

# 方法1：修改vgg16模型，在最后加一个线性层，即输入1000，输出10
vgg16_tr.classifier.add_module('add_linear',nn.Linear(1000,10))
print(vgg16_tr)

# 方法2：直接修改vgg16最后一个线性层输出为10
vgg16_fal.classifier[6] = nn.Linear(4096,10)
print(vgg16_fal)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27