📚 课程目录第2章模型量化

2. 模型量化

任务目标

本案例将训练一个简单的FP32 MNIST分类模型，使用TensorFlow Lite转换器执行全整型量化（Full Integer Quantization），将模型权重与激活从FP32转为INT8，并对比量化前后模型的推理精度，验证精度损失是否在可接受范围内。

环境与依赖

Python 3.8+
TensorFlow 2.12+（需包含TFLite Converter）
NumPy

完整实现代码

import tensorflow as tf
import numpy as np

# 1. 加载MNIST数据并训练简单FP32模型
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train, x_test = x_train.astype(np.float32) / 255.0, x_test.astype(np.float32) / 255.0
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

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.Flatten(),
    tf.keras.layers.Dense(10)
])
model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=['accuracy'])
model.fit(x_train, y_train, epochs=1, batch_size=128, validation_split=0.1, verbose=1)

# 2. 评估FP32模型精度
fp32_loss, fp32_acc = model.evaluate(x_test, y_test, verbose=0)
print(f"FP32 模型精度: {fp32_acc:.4f}")

# 3. 构建代表数据集（校准用，取100张测试图）
def representative_data_gen():
    for i in range(100):
        yield [x_test[i:i+1]]

# 4. 转换为全整型INT8量化模型
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_data_gen
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8
tflite_int8_model = converter.convert()

# 5. 保存量化模型
with open('mnist_int8.tflite', 'wb') as f:
    f.write(tflite_int8_model)

# 6. 加载TFLite模型并评估INT8精度
interpreter = tf.lite.Interpreter(model_content=tflite_int8_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']
output_scale, output_zero_point = output_details[0]['quantization']

# 对测试集进行INT8推理
correct = 0
for i in range(len(x_test)):
    # 输入量化
    input_data = x_test[i:i+1]
    input_data = input_data / input_scale + input_zero_point
    input_data = np.clip(input_data, -128, 127).astype(np.int8)
    
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()
    
    # 输出反量化
    output_data = interpreter.get_tensor(output_details[0]['index'])
    output_data = (output_data.astype(np.float32) - output_zero_point) * output_scale
    
    if np.argmax(output_data[0]) == y_test[i]:
        correct += 1

int8_acc = correct / len(x_test)
print(f"INT8 量化模型精度: {int8_acc:.4f}")
print(f"精度损失: {fp32_acc - int8_acc:.4f}")

运行方式

将上述代码保存为 quantize_mnist.py；
在终端执行 python quantize_mnist.py；
脚本会自动下载MNIST数据集（首次运行需联网），训练1个epoch后完成转换与评估。

验证方法与判定标准

观察输出：脚本会依次打印FP32精度、INT8精度和精度损失值。
通过条件：INT8模型精度不低于FP32精度的98%（即损失小于约0.02），且生成的 mnist_int8.tflite 文件大小明显小于原始Keras模型（通常压缩至约1/4）。
附加检查：可用 tf.lite.Interpreter 加载生成的tflite文件，确认输入输出类型为 int8（代码中已通过 inference_input_type 和 inference_output_type 强制指定）。

预期输出（基于代码逻辑）

FP32 模型精度: 0.9680
INT8 量化模型精度: 0.9665
精度损失: 0.0015

以上数值为示例，实际运行会因随机初始化略有波动，但精度损失通常小于0.01。

常见问题

转换报错：若提示缺少代表数据集，请确认 representative_data_gen 函数已正确传入且每次yield形状为 (1, 28, 28, 1)。
精度骤降：若INT8精度低于0.9，可增加代表数据集样本数（如改为500张），或训练更多epoch提升FP32基线精度。
输入输出类型不符：若推理时发现输入仍是float32，请检查是否已设置 converter.inference_input_type = tf.int8。