📚 课程目录第1章基于TensorFlow的MNIST模型训练与量化感知训练基础

1. 基于TensorFlow的MNIST模型训练与量化感知训练基础

任务目标

完成一个基于TensorFlow的MNIST手写数字识别模型，从浮点模型训练到量化感知训练（QAT）的完整流程，最终导出可供TFLite转换的量化模型。通过本案例掌握QAT的基本操作步骤，为后续嵌入式端部署打下基础。

环境依赖

Python 3.8+
TensorFlow 2.12+（含TensorFlow Model Optimization Toolkit）
NumPy

完整实现代码

import tensorflow as tf
import numpy as np
import tensorflow_model_optimization as tfmot

# 1. 加载MNIST数据集
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train.reshape(-1, 28, 28, 1).astype('float32') / 255.0
x_test = x_test.reshape(-1, 28, 28, 1).astype('float32') / 255.0

# 2. 构建浮点基线模型
def build_model():
    model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(8, 3, activation='relu', input_shape=(28, 28, 1)),
        tf.keras.layers.MaxPooling2D(2),
        tf.keras.layers.Conv2D(16, 3, activation='relu'),
        tf.keras.layers.MaxPooling2D(2),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    return model

# 3. 训练浮点基线模型
base_model = build_model()
base_model.compile(optimizer='adam',
                   loss='sparse_categorical_crossentropy',
                   metrics=['accuracy'])
base_model.fit(x_train, y_train, epochs=3, batch_size=128, validation_split=0.1, verbose=1)

# 4. 应用量化感知训练
quantize_model = tfmot.quantization.keras.quantize_model
qat_model = quantize_model(base_model)
qat_model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
qat_model.fit(x_train, y_train, epochs=2, batch_size=128, validation_split=0.1, verbose=1)

# 5. 评估QAT模型
test_loss, test_acc = qat_model.evaluate(x_test, y_test, verbose=0)
print(f'QAT模型测试准确率: {test_acc:.4f}')

# 6. 导出为TFLite格式（含量化信息）
converter = tf.lite.TFLiteConverter.from_keras_model(qat_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()

with open('mnist_qat.tflite', 'wb') as f:
    f.write(tflite_model)

# 7. 验证TFLite模型推理
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 取测试集前5个样本验证
for i in range(5):
    test_input = np.expand_dims(x_test[i], axis=0).astype('float32')
    interpreter.set_tensor(input_details[0]['index'], test_input)
    interpreter.invoke()
    output = interpreter.get_tensor(output_details[0]['index'])
    pred = np.argmax(output[0])
    print(f'样本{i+1} 预测: {pred}, 真实标签: {y_test[i]}')

运行方式

pip install tensorflow tensorflow-model-optimization numpy
python mnist_qat.py

验证方法与判定标准

训练收敛验证：训练过程中loss持续下降，最终QAT模型测试准确率应不低于95%（MNIST简单任务，浮点模型通常可达98%以上，QAT后下降不超过1~2个百分点）。
量化模型导出验证：运行结束后当前目录生成mnist_qat.tflite文件，文件大小应明显小于未量化的浮点模型（通常约1/4大小）。
推理一致性验证：控制台输出的5个样本预测结果与真实标签一致，说明TFLite量化模型推理功能正常。

关键说明

QAT的核心机制是在训练过程中模拟量化误差，使模型权重适应低精度表示，从而减少部署时的精度损失。
第6步中tf.lite.Optimize.DEFAULT会触发权重量化，而QAT模型中的伪量化节点会指导转换器生成更优的量化参数。
如需进一步压缩为int8全整数量化（嵌入式端常用），需额外提供代表性数据集进行校准，本案例为最小闭环暂不展开。

上图展示了本案例的完整流程与预期效果。QAT通过在训练中模拟量化噪声，使模型在低精度部署时保持较高准确率，是嵌入式端模型部署前的重要优化步骤。