深度学习pytorch：VGG网络模型的使用、修改及保存、添加线性层、修改网络输出

# 现有网络模型的使用及修改
import torchvision
from torch import nn
# 加载预训练网络模型
vgg16_true = torchvision.models.vgg16(pretrained=True)   # 下载网络模型

vgg16_false = torchvision.models.vgg16(pretrained=False)   # 只是加载网络模型
print((vgg16_true))

train_data = torchvision.datasets.CIFAR10(r"C:\Users\123\Desktop\python4.7\test03\data", train=False, transform=torchvision.transforms.ToTensor(),
                                       download=True)

# 在最后一层添加线性层，使输出的分类类别为10
vgg16_true.add_module("add_linear", nn.Linear(1000, 10))   # 输入1000，输出10
#  (add_linear): Linear(in_features=1000, out_features=10, bias=True)
# 在classifier最后一层加线性层
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)

# 模型的保存
import torch
import torchvision

vgg16 = torchvision.models.vgg16(pretrained=False)
# 两种方式保存模型
# 保存方法一(模型和参数一起保存)
torch.save(vgg16, "vgg16_method1.pth")

# 保存方法二(仅保存参数，官方推荐)
torch.save(vgg16.state_dict(), "vgg16_method2.pth")   # 将模型状态(参数)保存为字典形式



# 所保存模型的调用
import torch
# 方式一加载模型
import torchvision

model1 = torch.load("vgg16_method1.pth")
print(model1)

# 方式二加载模型
model2 = torch.load("vgg16_method2.pth")
print(model2)

# 若要使用方式二加载模型参数和结构
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict((torch.load("vgg16_method2.pth")))
print(vgg16)

#