📚 课程目录第2章 TensorFlow模型量化感知训练（QAT）入门

2. TensorFlow模型量化感知训练（QAT）入门

任务目标

使用TensorFlow 2.x对简单CNN模型进行量化感知训练（QAT），导出全整数量化模型（int8），并对比量化前后模型在MNIST测试集上的精度变化，验证QAT流程的有效性。

环境依赖

需要TensorFlow 2.15+（含tensorflow-model-optimization库），安装命令：

pip install tensorflow tensorflow-model-optimization

完整实现代码

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
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]

# 2. 构建简单CNN（浮点模型）
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)
    ])
    return model

# 3. 训练浮点基线模型
base_model = build_model()
base_model.compile(optimizer='adam',
                   loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                   metrics=['accuracy'])
base_model.fit(x_train, y_train, epochs=2, 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=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                  metrics=['accuracy'])
qat_model.fit(x_train, y_train, epochs=1, batch_size=128, validation_split=0.1, verbose=1)

# 5. 评估浮点与QAT模型精度
_, base_acc = base_model.evaluate(x_test, y_test, verbose=0)
_, qat_acc = qat_model.evaluate(x_test, y_test, verbose=0)
print(f'浮点模型精度: {base_acc:.4f}')
print(f'QAT模型精度: {qat_acc:.4f}')

# 6. 导出全整数量化模型（int8）
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)])
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8
tflite_model = converter.convert()

# 保存.tflite文件
with open('qat_mnist_int8.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()

# 获取输入量化参数
input_scale, input_zero_point = input_details[0]['quantization']
correct = 0
for i in range(len(x_test)):
    # 输入量化
    input_data = np.clip(x_test[i] / input_scale + input_zero_point, -128, 127).astype(np.int8)
    interpreter.set_tensor(input_details[0]['index'], input_data[None, ...])
    interpreter.invoke()
    output = interpreter.get_tensor(output_details[0]['index'])
    if np.argmax(output[0]) == y_test[i]:
        correct += 1

tflite_acc = correct / len(x_test)
print(f'量化模型精度: {tflite_acc:.4f}')
print(f'模型大小: {len(tflite_model)} bytes')

运行方式

将上述代码保存为qat_mnist.py，直接执行：

python qat_mnist.py

训练过程约需2-3分钟（CPU），最终生成qat_mnist_int8.tflite文件。

验证方法与判定标准

程序运行结束后，观察终端输出：

精度对比：QAT模型精度与浮点模型精度差距通常在1%以内（预期浮点约0.98，QAT约0.97-0.98）；量化模型精度与QAT模型精度差距应小于0.5%。
模型大小：量化模型文件大小约为浮点模型的1/4（预期约10-15KB）。
通过条件：量化模型精度不低于浮点模型精度的95%，且推理结果与代码逻辑一致。

关键说明

QAT在训练过程中模拟量化误差，使模型权重适应低精度表示，因此量化后精度损失远小于训练后直接量化（PTQ）。
导出int8模型必须提供representative_dataset（校准数据集），用于确定激活值的动态范围。
推理时输入需手动量化到int8范围，输出为int8格式，需按输出量化参数反量化后才能获得最终logits。