from model import LeNet5    # 引入我们写好的神经网络（也是我们要训练的网络）
from torch.nn import CrossEntropyLoss  # 引入损失函数（用来做分类损失最佳）
from torch.optim import Adam  # 损失函数优化器（其他损失函数的最小值：使用梯度下降法）
from torchvision.datasets import MNIST # 数据集
from torchvision.transforms import Compose, ToTensor  # 图像转换为张量
from torch.utils.data import DataLoader # 批次数据集（数据集不是一次训练，而是分成多个批次）
import torch  # 调用torch的基本函数
import os   # 路径与文件处理

class Trainer:
    def __init__(self):
        # 判定电脑安装GPU的环境：cuda,cnn,pytorch-
        self.CUDA = False # torch.cuda.is_available()   # 返回逻辑值：True：支持GPU，False不支持GPU
        # 准备训练需要的数据、损失函数，优化器，学习率
        self.lr = 0.0001
        self.m_file = "lenet5.pt"  # 模型的保存文件
        self.net = LeNet5() # 需要训练的网络
        if self.CUDA:
            self.net.cuda()   # 把net网络模型存储到GPU上面

        # 加载上一次训练的模型
        if os.path.exists(self.m_file):
            # 存在就加载
            print("加载模型中...")
            state = torch.load(self.m_file)  # 读取文件
            self.net.load_state_dict(state)
        else:
            print("模型存在，重头训练")    
        self.loss_f = CrossEntropyLoss()  # 损失函数
        self.optimizer = Adam(self.net.parameters(), lr=self.lr)  # 优化器


        # 数据集 - 训练集
        self.trans = Compose([ToTensor()])
        self.ds_train = MNIST(root="datasets", download=True, train=True, transform=self.trans)
        self.bt_train = DataLoader(self.ds_train, batch_size=1000, shuffle=True) 

        # 数据集 - 验证集
        self.ds_valid = MNIST(root="datasets", download=True, train=False, transform=self.trans)
        self.bt_valid = DataLoader(self.ds_valid, batch_size=1000, shuffle=False) #  shuffle=False是否随机打乱



    def train_one(self):
        # 训练一轮epoch，60批次batch，每个批次1000张录像
        # 循环训练每个批次
        batch = 1
        for x, y in self.bt_train:
            # print(F"\t|-第{batch:02d}批次训练。")
            if self.CUDA:
                x = x.cuda()
                y = y.cuda()
            y_ = self.net(x)  # 进行预测
            # 计算误差
            loss = self.loss_f(y_, y)   # 使用真实的标签与预测标签计算sunshi
            # 优化卷积核与全链接矩阵
            self.optimizer.zero_grad()
            loss.backward()   # 使用导数计算梯度
            self.optimizer.step()  # 更新梯度，优化网络模型
            batch+=1

    @torch.no_grad()   # 该函数中的运算都不会进行求导跟踪
    def valid(self):
        # 使用测试数据集验证 (准确率，损失值)
        all_num = 0.0 # 验证的样本数
        acc_num = 0.0 # 识别正确数量
        all_loss = 0.0 # 累加每个样本的损失
        for t_x, t_y in self.bt_valid:
            if self.CUDA:
                t_x = t_x.cuda()
                t_y = t_y.cuda()
            # 统计样本数
            all_num  += len(t_x)
            #  预测
            t_y_ = self.net(t_x)
            # 判定准确
            y_cls = torch.argmax(t_y_, dim=1)
            # 统计正确率
            acc_num  += (y_cls == t_y).float().sum()
            # 统计损失
            all_loss += self.loss_f(t_y_, t_y)
        print(F"\t|- 训练损失：{all_loss/all_num:8.6f}")
        print(F"\t|- 准确率：{acc_num * 100.0/all_num:5.2f}%")



    def train(self, epoch):
        # 训练指定的论数
        for e in range(epoch):
            print(F"第{e:04d}轮训练。")
            self.train_one()
            self.valid()
            # 保存训练的模型
            torch.save(self.net.state_dict(), self.m_file)

if __name__ == "__main__":
    trainer = Trainer()
    trainer.train(50) # 训练10轮
    # print(torch.cuda.is_available()) 

"""
    244K = 6 * 5 * 5 * 8 矩阵  
          16 * 5 * 5 * 8 
         120 * 5 * 5 * 8
             120 * 84* 8
              84 * 10*8
"""