4. 量化感知训练实战

任务目标

在TinyML部署中，将浮点模型直接转换为int8量化模型通常会导致1%~3%的精度损失。本案例通过量化感知训练（QAT），在训练过程中模拟量化误差，使模型权重和激活值适应低精度表示，最终导出int8模型时精度接近原始浮点模型。目标：使用TensorFlow/Keras对MNIST分类模型执行QAT，对比浮点模型与QAT后int8模型的精度差异。

环境依赖

Python 3.9+，TensorFlow 2.13+（含tfmot），NumPy。安装命令：

pip install tensorflow==2.13.* tensorflow-model-optimization==0.7.* numpy

完整实现

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

# 1. 加载MNIST数据
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

# 2. 构建浮点基线模型
def build_model():
    model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=(28, 28)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(128, activation='relu'),
        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. 应用QAT（在浮点模型基础上）
qat_model = tfmot.quantization.keras.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模型
float_acc = base_model.evaluate(x_test, y_test, verbose=0)[1]
qat_acc = qat_model.evaluate(x_test, y_test, verbose=0)[1]
print(f"浮点模型精度: {float_acc:.4f}")
print(f"QAT模型精度: {qat_acc:.4f}")

# 6. 导出int8 TFLite模型（QAT）
converter = tf.lite.TFLiteConverter.from_keras_model(qat_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = lambda: ([x_test[i:i+1] for i in range(100)])
int8_model = converter.convert()
with open('qat_mnist.tflite', 'wb') as f:
    f.write(int8_model)

# 7. 验证int8 TFLite模型精度
interpreter = tf.lite.Interpreter(model_content=int8_model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

correct = 0
for i in range(len(x_test)):
    inp = np.expand_dims(x_test[i].astype(np.float32), axis=0)
    interpreter.set_tensor(input_details[0]['index'], inp)
    interpreter.invoke()
    pred = np.argmax(interpreter.get_tensor(output_details[0]['index']))
    if pred == y_test[i]:
        correct += 1
int8_acc = correct / len(x_test)
print(f"int8 TFLite模型精度: {int8_acc:.4f}")

运行与验证

执行方式：将上述代码保存为qat_mnist.py，在终端运行python qat_mnist.py。训练过程约需2~3分钟（CPU）。

通过条件：

• 程序正常结束，无报错；
• 输出三个精度值，且满足：int8 TFLite模型精度与浮点模型精度差距小于1.5个百分点；
• 生成qat_mnist.tflite文件（约25KB）。

预期结果参考：浮点模型精度约0.97~0.98，QAT模型精度约0.97，int8模型精度约0.96~0.97。若直接使用训练后量化（非QAT），int8精度通常降至0.95以下，本案例通过QAT有效缩小了这一差距。

关键说明

• 第4步quantize_model在浮点模型结构上插入伪量化节点，训练时这些节点模拟int8舍入误差，使权重适应量化噪声；
• 第6步导出时无需再提供代表性数据集做校准（QAT已内化量化范围），但代码中仍保留以兼容旧版API；
• 若需部署到微控制器，可将生成的.tflite文件转换为C数组（使用xxd -i），此处不展开。