import torch
from torch import nn  # 导入神经网络模块
from torchvision import transforms  # 导入数据集模块
from torchvision import datasets  # 导入数据集模块
from torch.utils.data import DataLoader  # 导入数据加载器
import torch.nn.functional as F  # 导入函数式接口（如ReLU、sigmoid等）
import torch.optim as optim  # 导入优化器模块

batch_size = 64  # 设置批次大小，每次训练使用64个样本
# 定义数据预处理管道
transform = transforms.Compose([transforms.ToTensor(),  # 将PIL图像或numpy数组转为Tensor，并归一化到[0,1]
                                transforms.Normalize((0.1307,), (0.3081,))])  # 标准化：(x-mean)/std
# 加载MNIST训练集
train_dataset = datasets.MNIST(root='../dataset/mnist/', # 数据集存储路径
                               train=True,  # 使用训练集
                               transform=transform,  # 应用上面定义的数据变换
                               download=True)  # 如果本地没有则下载
# 创建训练数据加载器
train_loader = DataLoader(train_dataset,  # 训练数据集
                          batch_size=batch_size,  # 批次大小64
                          shuffle=True)  # 每个epoch打乱数据顺序，防止过拟合
# 加载MNIST测试集
test_dataset = datasets.MNIST(root='../dataset/mnist/',  # 同训练集路径
                              train=False,  # 使用测试集
                              transform=transform,  # 应用相同的变换
                              download=True)  # 如果本地没有则下载
test_loader = DataLoader(test_dataset,  # 测试数据集
                         batch_size=batch_size,  # 批次大小64
                         shuffle=False)  # 测试集不需要打乱，保持固定顺序

class Net(torch.nn.Module):  # 继承nn.Module基类
    def __init__(self):
        super(Net, self).__init__()  # 调用父类初始化
        self.l1 = torch.nn.Linear(784, 512)  # 输入784维，输出512维
        self.l2 = torch.nn.Linear(512, 256)  # 输入512维，输出256维
        self.l3 = torch.nn.Linear(256, 128)  # 输入256维，输出128维
        self.l4 = torch.nn.Linear(128, 64)   # 输入128维，输出64维
        self.l5 = torch.nn.Linear(64, 10)    # 输入64维，输出10维（10个数字类别）

    def forward(self, x):  # 定义前向传播
        x = x.view(-1, 784)  # ☆将图像展平，-1自动计算批次大小：[batch, 1, 28, 28] → [batch, 784]
        x = F.relu(self.l1(x))  # 第一层：线性变换 + ReLU激活
        x = F.relu(self.l2(x))  # 第二层：线性变换 + ReLU激活
        x = F.relu(self.l3(x))  # 第三层：线性变换 + ReLU激活
        x = F.relu(self.l4(x))  # 第四层：线性变换 + ReLU激活
        return self.l5(x)  # 第五层：线性变换（不加激活函数）

model = Net()  # 实例化神经网络模型
criterion = torch.nn.CrossEntropyLoss()  # 定义损失函数（适用于多分类问题，内部包含LogSoftmax）
optimizer = torch.optim.SGD(model.parameters(),  # 需要优化的参数（模型所有权重和偏置）
                            lr=0.01,  # 学习率
                            momentum=0.5)  # 动量，加速收敛并减少震荡

def train(epoch):
    running_loss = 0.0  # 累计损失，用于计算平均损失
    # 遍历所有训练批次
    for batch_idx, data in enumerate(train_loader, 0):
        inputs, target = data  # 解包：inputs是图像，target是标签
        outputs = model(inputs)  # 前向传播：输入→模型→输出
        loss = criterion(outputs, target)  # 计算损失：预测值 vs 真实值
        optimizer.zero_grad()  # 清空梯度缓存，防止梯度累积
        loss.backward()  # 反向传播：计算所有参数的梯度
        optimizer.step()  # 参数更新：根据梯度更新模型参数
        running_loss += loss.item()  # 累加当前批次损失
        if batch_idx % 300 == 299:  # 每300个批次打印一次平均损失
            print('[%d, %5d] loss: %.3f' % (epoch + 1, batch_idx + 1, running_loss / 300))
            running_loss = 0.0  # 重置累计损失

def test():
    correct = 0  # 正确预测的样本数
    total = 0    # 总样本数
    with torch.no_grad():  # 测试时不计算梯度，节省内存和计算
        for data in test_loader:  # 遍历所有测试批次
            images, labels = data  # 解包
            outputs = model(images)  # 前向传播获取预测值
            # torch.max返回两个值：最大值和最大值索引
            # dim=1表示在类别维度（第1维）找最大值
            # predicted是预测的类别索引（0-9）
            _, predicted = torch.max(outputs.data, dim=1)
            total += labels.size(0)  # 累加当前批次样本数
            correct += (predicted == labels).sum().item()  # 累加正确预测数
    # 计算并打印准确率
    print('Accuracy of the test set: %d %%' % (100 * correct / total))

if __name__ == '__main__':  # 主程序入口，防止多进程问题
    for epoch in range(10):  # 训练10个epoch
        train(epoch)  # 训练一个epoch
        test()        # 测试当前模型性能