#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@Time        : 2022/5/8 08:19
@Author      : Albert Darren
@Contact     : 2563491540@qq.com
@File        : train.py
@Version     : Version 1.0.0
@Description : TODO
@Created By  : PyCharm
"""
from torchvision.datasets import CIFAR10
from torchvision import transforms
from torch.utils.data import DataLoader
from model import CF10Net
from torch.nn import CrossEntropyLoss
from torch.optim import SGD
from torch.utils.tensorboard import SummaryWriter
from torch import save, no_grad

# 准备数据集
train_dataset = CIFAR10(root="cifai10", train=True, transform=transforms.ToTensor(), download=True)
test_dataset = CIFAR10(root="cifai10", train=False, transform=transforms.ToTensor(), download=True)
# 数据探索
train_dataset_size = len(train_dataset)
test_dataset_size = len(test_dataset)
# 数据加载
train_dataloader = DataLoader(dataset=train_dataset, batch_size=64)
test_dataloader = DataLoader(dataset=test_dataset, batch_size=64)
# 创建神经网络
cf10net = CF10Net()
# 损失函数
loss_cross = CrossEntropyLoss()
# 学习速率
learning_rate = 1e-2
# 优化器
optimizer = SGD(cf10net.parameters(), lr=learning_rate)
# 设置训练轮数
epoch = 10
# 事件写入对象
writer = SummaryWriter(log_dir="cf10net_logs")
for i in range(epoch):
    print("第{}轮训练开始....".format(i + 1))
    total_train_acc = 0
    total_test_acc = 0
    # 训练步骤开始
    cf10net.train()
    for data in train_dataloader:
        images, targets = data
        # zeros all parameters gradients
        optimizer.zero_grad()
        # compute output
        outputs = cf10net(images)
        total_train_acc += (outputs.argmax(dim=1) == targets).sum().item()
        # get loss
        loss = loss_cross(outputs, targets)
        # backward
        loss.backward()
        # update weights
        optimizer.step()
    # 测试步骤开始
    cf10net.eval()
    with no_grad():
        for data in test_dataloader:
            images, targets = data
            outputs = cf10net(images)
            total_test_acc += (outputs.argmax(dim=1) == targets).sum().item()
    writer.add_scalars("ACC", {"train dataset acc": total_train_acc / train_dataset_size,
                               "test dataset acc": total_test_acc / test_dataset_size}, i + 1)
    save(cf10net.state_dict(), "models/cf10net_{}.pth".format(i + 1))
    print("模型保存成功！")
writer.close()