#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@Time        : 2022/5/11 15:18
@Author      : Albert Darren
@Contact     : 2563491540@qq.com
@File        : train_gpu_2.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 torch.nn import CrossEntropyLoss
from torch.optim import SGD
from torch.utils.tensorboard import SummaryWriter
from torch import save, no_grad, cuda, device
from model import CF10Net

# 定义运算的设备
train_device = device("cuda" if cuda.is_available() else "cpu")
# 准备数据集
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()
cf10net = cf10net.to(train_device)
# 损失函数
loss_cross = CrossEntropyLoss()
loss_cross = loss_cross.to(train_device)
# 学习速率
learning_rate = 1e-3
# 优化器
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
        # 将训练数据样本及其标签迁移到运算设备
        images = images.to(train_device)
        targets = targets.to(train_device)
        # 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
            # 将测试数据样本及其标签迁移到运算设备
            images = images.to(train_device)
            targets = targets.to(train_device)
            outputs = cf10net(images)
            total_test_acc += (outputs.argmax(dim=1) == targets).sum().item()
    writer.add_scalars("GPU2 ACC", {"train dataset gpu2 acc": total_train_acc / train_dataset_size,
                                   "test dataset gpu2 acc": total_test_acc / test_dataset_size}, i + 1)
    save(cf10net.state_dict(), "models/cf10net_gpu2_{}.pth".format(i + 1))
    print("模型保存成功！")
writer.close()