# LeNet


import d2lzh as d2l
import mxnet as mx
from mxnet import autograd, gluon, init, nd
from mxnet.gluon import loss as gloss, nn
import time

# 构建模型

# 输入数据：神经网络的输入是一个4维张量，形状为(批量大小, 通道数, 高度, 宽度)。这里的通道数表示图像的颜色通道数，对于灰度图像为1，对于RGB图像为3。
#
# 第1个卷积层：输入的4维张量通过第一个卷积层。这个卷积层会对输入进行卷积操作，得到6个输出通道的特征图。使用sigmoid激活函数处理特征图。
#
# 第1个最大池化层：经过第1个卷积层的输出被送入第1个最大池化层。最大池化层将特征图的空间维度（高度和宽度）减小一半，并保留每个2x2区域内的最大值。
#
# 第2个卷积层：第1个最大池化层的输出被送入第2个卷积层。这个卷积层将6个输入通道的特征图转换为16个输出通道的特征图。同样使用sigmoid激活函数。
#
# 第2个最大池化层：经过第2个卷积层的输出被送入第2个最大池化层。这个池化层也会将特征图的空间维度减小一半。
#
# 全连接层1：第2个最大池化层的输出被展平为一个1维张量，并送入全连接层1。这个全连接层有120个神经元，使用sigmoid激活函数。
#
# 全连接层2：全连接层1的输出被送入全连接层2。这个全连接层有84个神经元，同样使用sigmoid激活函数。
#
# 输出层：全连接层2的输出被送入输出层。这是一个没有显式激活函数的全连接层，通常称为“输出层”。它有10个神经元，对应于最终的分类类别数量。
#
# 输出结果：输出层的结果就是最终的预测结果，其中每个神经元对应一个类别的预测概率。通常，我们会选取概率最高的类别作为最终的预测类别。


net = nn.Sequential()  # 创建一个空的神经网络容器
net.add(nn.Conv2D(channels=6, kernel_size=5, activation='sigmoid'),
        # 添加第一个卷积层，有6个输出通道，卷积核大小为5x5，使用sigmoid激活函数
        nn.MaxPool2D(pool_size=2, strides=2),
        # 添加第一个最大池化层，池化窗口大小为2x2，步幅为2
        nn.Conv2D(channels=16, kernel_size=5, activation='sigmoid'),
        # 添加第二个卷积层，有16个输出通道，卷积核大小为5x5，使用sigmoid激活函数
        nn.MaxPool2D(pool_size=2, strides=2),
        # 添加第二个最大池化层，池化窗口大小为2x2，步幅为2
        # Dense会默认将(批量大小, 通道, 高, 宽)形状的输入转换成
        # (批量大小, 通道 * 高 * 宽)形状的输入
        nn.Dense(120, activation='sigmoid'),
        # 添加一个全连接层，有120个神经元，使用sigmoid激活函数.
        # 进入这个全连接层之前，会被展平，也就是从一个4维张量（batch_size, channels, height, width）
        # 转换为一个2维张量（batch_size, channels * height * width）。激活函数使用sigmoid。
        nn.Dense(84, activation='sigmoid'),
        # 添加另一个全连接层，有84个神经元，使用sigmoid激活函数
        nn.Dense(10))  # 最后添加一个输出层，有10个神经元，用于分类任务

ctx = d2l.try_gpu()

# 构造样本

# 构造一个高和宽均为28的单通道数据样本，并逐层进行前向计算来查看每个层的输出形状。

# 创建输入数据

x = nd.random_uniform(shape=(1, 1, 28, 28))  # 创建一个形状为(1, 1, 28, 28)的随机输入张量

net.initialize()  # 初始化神经网络

# 对每一层进行前向计算，并打印输出形状

for layer in net:  # 对神经网络中的每一层进行循环遍历
    x = layer(x)  # 将输入数据通过当前层进行前向计算
    print(layer.name, 'output shape:\t', x.shape)  # 打印当前层的名称和输出形状

# 获取数据和训练模型

batch_size = 256  # 定义批量大小
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size=batch_size)  # 加载Fashion-MNIST数据集


# 计算精确率
def evaluate_accuracy(data_iter, net):  # 定义计算精确率的函数
    acc_sum, n = 0.0, 0  # 初始化正确分类数和样本总数

    for x, y in data_iter:  # 对数据迭代器进行遍历，获取输入和标签
        y = y.astype('float32')  # 将标签转换为浮点数类型
        acc_sum += (net(x).argmax(axis=1) == y).sum().asscalar()  # 统计正确分类的个数
        n += y.size  # 统计样本总数
    return acc_sum / n  # 返回精确率（正确分类数除以样本总数）


# 训练模型
def train(net, train_iter, test_iter, batch_size, trainer, num_epochs):  # 定义训练模型的函数
    loss = gloss.SoftmaxCrossEntropyLoss()  # 定义交叉熵损失函数

    for epoch in range(num_epochs):  # 对每个epoch进行循环
        train_loss, train_acc, n, start = 0.0, 0.0, 0, time.time()  # 初始化训练损失、训练精确率、样本数和起始时间
        for data, label in train_iter:  # 对训练数据进行迭代
            data, label = data.as_in_context(ctx), label.as_in_context(ctx)
            with autograd.record():  # 开启自动求导
                output = net(data)  # 进行前向计算，获取模型输出
                _loss = loss(output, label).sum()  # 计算损失

            _loss.backward()  # 反向传播，计算梯度
            trainer.step(batch_size)  # 更新模型参数
            label = label.astype('float32')  # 将标签转换为浮点数类型

            train_loss += _loss.asscalar()  # 累加训练损失
            train_acc += (output.argmax(axis=1) == label).sum().asscalar()  # 累加训练精确率

            n += label.size  # 累加样本数
        test_acc = evaluate_accuracy(test_iter, net)  # 计算测试精确率

        print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f, ''time %.1f sec' % (
            epoch + 1, train_loss / n, train_acc / n, test_acc, time.time() - start))  # 打印每个epoch的训练损失、训练精确率、测试精确率和训练时间


# 训练模型

lr, num_epochs = 0.9, 5  # 定义学习率和训练的epoch数
net.initialize(force_reinit=True, init=init.Xavier())  # 强制初始化神经网络的参数，使用Xavier初始化方法

trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': lr})  # 定义优化器
train(net, train_iter, test_iter, batch_size, trainer, num_epochs)  # 训练模型，获取训练和测试精确率