# 导入必要的库和模块
import tensorflow as tf  # 引入TensorFlow库,用于构建和训练神经网络模型
from tensorflow.keras.datasets import mnist  # 引入MNIST手写数字数据集
from tensorflow.keras.layers import Flatten, Dense  # 引入神经网络层:Flatten和Dense
from tensorflow.keras.models import Sequential  # 引入Sequential模型
from tensorflow.keras.utils import to_categorical  # 引入to_categorical函数,用于将标签转换为one-hot编码
import numpy as np  # 引入NumPy库,用于数组操作和数值计算

# 1. 引入数据,对数据进行预处理
(X_train, y_train), (X_test, y_test) = mnist.load_data()  # 加载MNIST数据集,并将其拆分为训练集和测试集
X_train = X_train / 255.0  # 将训练集像素值归一化到[0,1]范围内,提高模型训练效果
X_test = X_test / 255.0  # 将测试集像素值归一化到[0,1]范围内,提高模型推理效果
y_train = to_categorical(y_train)  # 将训练集标签转换为one-hot编码,便于模型训练
y_test = to_categorical(y_test)  # 将测试集标签转换为one-hot编码,便于模型评估

# 2. 构建网络模型
model = Sequential()  # 创建Sequential模型
model.add(Flatten(input_shape=(28, 28)))  # 添加Flatten层,将二维图像数据展平为一维向量
model.add(Dense(20, activation=tf.nn.relu))  # 添加全连接层,使用ReLU激活函数,该层有20个神经元
model.add(Dense(10, activation=tf.nn.softmax))  # 添加输出层,使用Softmax激活函数,该层有10个神经元,对应10个类别

# 3. 编译模型,指定优化器、损失函数和评估指标
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# 4. 训练模型
model.fit(X_train, y_train, epochs=20, validation_split=0.2)  # 训练模型,设置训练轮数为20,验证集比例为20%

# 5. 评估模型
test_loss, test_acc = model.evaluate(X_test, y_test)  # 在测试集上评估模型性能
print('Test accuracy:', test_acc)  # 输出测试集准确率

# 6. 模型推理
X_infer = np.array([X_test[0], X_test[1]])  # 从测试集中选择前两个样本作为推理数据
y_infer = model.predict(X_infer)  # 使用训练好的模型对推理数据进行预测

# 7.保存模型到硬盘
model_path = './mnist_model'  # 指定保存模型的路径和文件名,便于后续加载
tf.keras.models.save_model(model, model_path)  # 使用save_model函数将训练好的模型保存到指定路径,实现模型持久化