"""
    使用官方提供的网络结构
        # 作为前置网络，按需修改结构
"""
import torch
import torchvision
from torch import nn

from torch.nn import Conv2d, MaxPool2d, ReLU, Sigmoid, Linear, Flatten, Sequential
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter("logs/018")

dataset = torchvision.datasets.CIFAR10(root="./visionData", train=False, transform=torchvision.transforms.ToTensor(),
                                       download=True)
dataloader = DataLoader(dataset, batch_size=1)

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

vgg16_false = torchvision.models.vgg16(pretrained=False)  # 未训练
vgg16_true = torchvision.models.vgg16(pretrained=True)  # 训练好的网络
print(vgg16_true)

# 给现有网络结构添加层
vgg16_true.add_module("add_linear", nn.Linear(1000, 10))
print(vgg16_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)

writer.close()