现有网络模型的使用及修改

gxyqn626

于 2023-02-22 15:23:36 发布

阅读量69

点赞数

文章标签：深度学习 Powered by 金山文档

本文链接：https://blog.csdn.net/gxyqn626/article/details/129157111

版权

该代码示例导入了torch和torchvision库，创建了两个VGG16模型，一个预训练（在ImageNet上），另一个未预训练。预训练模型加载了在ImageNet数据集上的权重。然后，对预训练模型的分类层进行了调整，以适应CIFAR10数据集的10个类别。未预训练的模型手动修改了最后一层的线性层以匹配CIFAR10的输出类别。

摘要由CSDN通过智能技术生成

pretrained(bool)：如果是True,说明这个模型在ImageNet数据集上已经训练好了。

progress(bool)：如果为True,显示一个下载的进度条。

import torch
import torchvision
import torchvision.models as models
# train_data = torchvision.datasets.ImageNet("../data_image_net",split='train', download=True,transform=torchvision.transforms.ToTensor())
from torch import nn

vgg16_false = models.vgg16(pretrained=False) #参数是默认的参数，是初始化的
vgg16_true = models.vgg16(pretrained=True) # 要去下载在ImageNet中训练好的参数
print("ok")
print(vgg16_true)

train_data = torchvision.datasets.CIFAR10('../data',train=True,transform=torchvision.transforms.ToTensor(),download=True)
vgg16_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)