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)
pytorch入门17:现有模型的使用和修改
最新推荐文章于 2023-08-21 16:54:49 发布