Pytorch学习日记——现有网络模型的使用及修改

学习视频——B站【小土堆】

代码

import torchvision
from torch import nn

# train_data = torchvision.datasets.ImageNet("..data_image_net", split="train", download=True, transform=torchvision.transforms.ToTensor)

vgg16_false = torchvision.models.vgg16(pretrained=False)
vgg16_true = torchvision.models.vgg16(pretrained=True)
# print("ok")

print(vgg16_true)

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

# vgg16_true.add_module("add_linear", nn.Linear(1000, 10)) # 在最后增加一个名为add_linear的线性层

vgg16_true.classifier.add_module("add_linear", nn.Linear(1000, 10)) # 在classifier最后增加一个名为add_linear的线性层
print(vgg16_true)

print(vgg16_false)
vgg16_false.classifier[6] = nn.Linear(4096, 10) # 将序号6的线性层改为10输出
print(vgg16_false)