"""
    模型保存和加载（两种方式）
        # 保存完整网络模型
        # 保存网络模型参数
"""
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/019")

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

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

# 保存方式1：直接保存完整模型
print(vgg16_false)
torch.save(vgg16_false, "./model_save/vgg16_method1.pth")

# 对应加载方法
model_load1 = torch.load("./model_save/vgg16_method1.pth")  # 加载模型
print(model_load1)

# 保存方式2：保存模型参数（官方推荐，占用空间小）
path = "./model_save/vgg16_method2.pth"
torch.save(vgg16_false.state_dict(), path)

# 对应加载方法
new_vgg16 = torchvision.models.vgg16(pretrained=False)  # 获取模型
model_param = torch.load(path)  # 获取参数
new_vgg16.load_state_dict(model_param)  # 装载参数
print(new_vgg16)


# 陷阱
class Gisleung(nn.Module):
    def __init__(self):
        super(Gisleung, self).__init__()
        self.conv1 = nn.Conv2d(3, 64)

    def forward(self, x):
        x = self.conv1(x)
        return x


# 方式一保存
path = "./model_save/gisleung_method2.pth"
gis = Gisleung()
torch.save(gis, path)

# 方式一加载时候，需要在当前程序中实现定义Gisleung
model = torch.load(path)
print(model)

writer.close()