pytorch现有模型的使用和修改

最新推荐文章于 2024-08-10 23:05:37 发布

原创最新推荐文章于 2024-08-10 23:05:37 发布 · 264 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能 #python

该篇博客介绍了如何使用预训练的VGG16模型，分别加载了预训练权重为True和False的情况，并在CIFAR10数据集上进行了操作。对于预训练为True的模型，添加了一个新的线性层以适应10类的输出；而对于预训练为False的模型，直接修改了最后一层的全连接层来匹配分类任务。

#model_pretrained


VGG16
vgg16_false = torchvision.models.vgg16(pretrained = False)
vgg16_true = torchvision.models.vgg16(pretrained = True)
print(vgg16_true)


train_data = torchvision.datasets.CIFAR10("./data",train=True, transform=torchvision.transforms.ToTensor(),download=Ture)
#增加线性层
vgg10_true.classifier.add_module("add_linear",nn.Linear(1000,10))
print(vgg16_true)

print(vgg16_false)
vgg16_false.classifier[6] = nn.Linear(4096,10)
print(vgg16_false)