#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@Time        : 2022/5/10 08:17
@Author      : Albert Darren
@Contact     : 2563491540@qq.com
@File        : train_gpu_1.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,cuda

# 准备数据集
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()
if cuda.is_available():
    # 将模型参数及缓冲移动到GPU，注意时机为新建模型对象之后
    cf10net = cf10net.cuda()
# 损失函数
loss_cross = CrossEntropyLoss()
# 将模型参数及缓冲移动到GPU，注意时机为新建损失函数对象之后
loss_cross = loss_cross.cuda()
# 学习速率
learning_rate = 1e-2
# 优化器
optimizer = SGD(cf10net.parameters(), lr=learning_rate)
# 设置训练轮数
epoch = 20
# 事件写入对象
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
        if cuda.is_available():
            # 将批采样样本及其标签
            images = images.cuda()
            targets = targets.cuda()
        # 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
            if cuda.is_available():
                # 将批采样样本及其标签
                images = images.cuda()
                targets = targets.cuda()
            outputs = cf10net(images)
            total_test_acc += (outputs.argmax(dim=1) == targets).sum().item()
    writer.add_scalars("GPU1 ACC", {"train dataset gpu1 acc": total_train_acc / train_dataset_size,
                                   "test dataset gpu1 acc": total_test_acc / test_dataset_size}, i + 1)
    save(cf10net.state_dict(), "models/cf10net_gpu1_{}.pth".format(i + 1))
    print("模型保存成功！")
writer.close()