# Author: Liuxin
# Time: 2022/4/21
# 导入训练所需包
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
import torch.nn.functional as F
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

# 数据读取
input_size = 28
num_classes = 10
num_epochs = 3
batch_size = 80

# 加载数据集，第一次加载时需要进行下载
train_dataset = datasets.MNIST(root='./data',
                               train=True,
                               transform=transforms.ToTensor(),
                               download=False)
test_dataset = datasets.MNIST(root='./data',
                              train=False,
                              transform=transforms.ToTensor(),
                              download=False)
# 构建batch数据
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          shuffle=True)


# 设计网络结构
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(  # 原始输入：28*28*1
            nn.Conv2d(in_channels=1,  # 灰度图
                      out_channels=16,  # 输出的特征图个数，即采用16个卷积核，生成16个特征图
                      kernel_size=5,  # 核函数大小
                      stride=1,  # 步长
                      padding=2),  # 如果希望卷积后和原来图像大小一样，需设置padding=(kernal_size-1)/2 if stride=1
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)  # 输出尺寸：14*14*16
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=16,
                      out_channels=32,
                      kernel_size=5,
                      stride=1,
                      padding=2),  # 输出尺寸：14*14*32
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)  # 输出尺寸：7*7*32
        )
        self.out = nn.Linear(32 * 7 * 7, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0),
                   -1)  # 展平层，x.size(0)指batchsize的值，最后通过x.view(x.size(0), -1)将tensor的结构转换为了(batchsize, channels*x*y)，即为(batch_size,32*7*7)
        out = self.out(x)
        return out


# 准确率作为评估标准
def accuracy(predictions, labels):
    pred = torch.argmax(predictions, 1)  # 获取最大值位置的索引，即预测类别
    rights = torch.sum(pred == labels.data)  # 或torch.sum(pred.eq(labels))
    return rights, len(labels)


# 训练网络模型
net = CNN()  # 实例化
criterion = nn.CrossEntropyLoss()  # 交叉熵损失函数
optimizer = optim.Adam(net.parameters(), lr=0.001)  # 优化器
# 开始训练
for epoch in range(num_epochs):
    train_rights = []
    for idx, (data, target) in enumerate(train_loader):
        net.train()  # 进行训练
        output = net(data)  # 预测输出
        loss = criterion(output, target)  # 计算损失
        optimizer.zero_grad()  # 梯度置0
        loss.backward()  # 反向传播
        optimizer.step()  # 梯度更新
        right = accuracy(output, target)  # 计算准确率
        train_rights.append(right)

        if idx % 100 == 0:
            net.eval()  # 进行测试
            test_rights = []
            for (test_data, test_target) in test_loader:
                test_output = net(test_data)
                rights = accuracy(output, test_target)  # 返回2个值
                test_rights.append(rights)

            # 计算准确率
            train_r = (sum([tup[0] for tup in train_rights]),
                       sum([tup[1] for tup in train_rights]))
            test_r = (sum([tup[0] for tup in test_rights]),
                      sum([tup[1] for tup in test_rights]))
            print('当前epoch：{}[{}/{} ({:.2f})%]\t 损失：{:.6f}\t 训练准确率：{:.2f}%\t 测试准确率：{:.2f}%'.format(
                epoch,
                idx * batch_size, len(train_loader.dataset),
                100 * idx / len(train_loader),
                loss.data,
                100 * train_r[0] / train_r[1],
                100 * test_r[0] / test_r[1]
            ))
#结果
"""
当前epoch：0[0/60000 (0.00)%]	 损失：2.306161	 训练准确率：7.50%	 测试准确率：9.76%
当前epoch：0[8000/60000 (13.33)%]	 损失：0.441864	 训练准确率：75.31%	 测试准确率：11.00%
当前epoch：0[16000/60000 (26.67)%]	 损失：0.241250	 训练准确率：84.07%	 测试准确率：9.89%
当前epoch：0[24000/60000 (40.00)%]	 损失：0.166731	 训练准确率：87.62%	 测试准确率：9.89%
当前epoch：0[32000/60000 (53.33)%]	 损失：0.064496	 训练准确率：89.68%	 测试准确率：10.15%
当前epoch：0[40000/60000 (66.67)%]	 损失：0.193250	 训练准确率：91.05%	 测试准确率：9.85%
当前epoch：0[48000/60000 (80.00)%]	 损失：0.049698	 训练准确率：92.12%	 测试准确率：9.76%
当前epoch：0[56000/60000 (93.33)%]	 损失：0.017600	 训练准确率：92.91%	 测试准确率：10.09%
当前epoch：1[0/60000 (0.00)%]	 损失：0.066587	 训练准确率：98.75%	 测试准确率：10.13%
当前epoch：1[8000/60000 (13.33)%]	 损失：0.024357	 训练准确率：98.00%	 测试准确率：9.95%
当前epoch：1[16000/60000 (26.67)%]	 损失：0.140894	 训练准确率：97.92%	 测试准确率：9.79%
当前epoch：1[24000/60000 (40.00)%]	 损失：0.056632	 训练准确率：97.93%	 测试准确率：10.18%
当前epoch：1[32000/60000 (53.33)%]	 损失：0.249817	 训练准确率：97.94%	 测试准确率：10.23%
当前epoch：1[40000/60000 (66.67)%]	 损失：0.095753	 训练准确率：97.96%	 测试准确率：9.76%
当前epoch：1[48000/60000 (80.00)%]	 损失：0.084824	 训练准确率：98.01%	 测试准确率：10.19%
当前epoch：1[56000/60000 (93.33)%]	 损失：0.046403	 训练准确率：98.08%	 测试准确率：9.66%
当前epoch：2[0/60000 (0.00)%]	 损失：0.004661	 训练准确率：100.00%	 测试准确率：9.92%
当前epoch：2[8000/60000 (13.33)%]	 损失：0.021806	 训练准确率：98.68%	 测试准确率：9.68%
当前epoch：2[16000/60000 (26.67)%]	 损失：0.053609	 训练准确率：98.58%	 测试准确率：10.17%
当前epoch：2[24000/60000 (40.00)%]	 损失：0.127273	 训练准确率：98.58%	 测试准确率：9.47%
当前epoch：2[32000/60000 (53.33)%]	 损失：0.014368	 训练准确率：98.61%	 测试准确率：9.74%
当前epoch：2[40000/60000 (66.67)%]	 损失：0.018668	 训练准确率：98.59%	 测试准确率：10.52%
当前epoch：2[48000/60000 (80.00)%]	 损失：0.047745	 训练准确率：98.60%	 测试准确率：10.19%
当前epoch：2[56000/60000 (93.33)%]	 损失：0.063654	 训练准确率：98.63%	 测试准确率：9.91%
"""