📚 课程目录第1章基于TensorFlow的MNIST模型量化

1. 基于TensorFlow的MNIST模型量化

任务目标

使用TensorFlow 2.x训练一个MNIST手写数字识别模型，通过训练后量化（Post-Training Quantization）将float32模型转换为int8模型，并对比量化前后的模型文件大小与推理精度变化。

环境依赖

Python 3.8+
TensorFlow 2.13+（需包含量化相关API）
NumPy

完整实现代码

import tensorflow as tf
import numpy as np
import os
import tempfile

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

# 2. 构建并训练模型
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(8, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(x_train, y_train, epochs=2, batch_size=128, validation_split=0.1, verbose=1)

# 3. 评估原始模型
_, float_acc = model.evaluate(x_test, y_test, verbose=0)
print(f"Float32模型精度: {float_acc:.4f}")

# 4. 保存原始模型
float_model_path = "/tmp/mnist_float_model.h5"
model.save(float_model_path)
float_size = os.path.getsize(float_model_path) / 1024.0
print(f"Float32模型大小: {float_size:.2f} KB")

# 5. 训练后量化（int8）
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = lambda: (
    tf.data.Dataset.from_tensor_slices((x_train[:100], y_train[:100]))
    .batch(1)
    .map(lambda x, y: x)
    .take(100)
    .as_numpy_iterator()
)
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()

# 6. 保存量化模型
quant_model_path = "/tmp/mnist_quant_model.tflite"
with open(quant_model_path, "wb") as f:
    f.write(tflite_model)
quant_size = os.path.getsize(quant_model_path) / 1024.0
print(f"Int8量化模型大小: {quant_size:.2f} KB")
print(f"模型压缩比: {float_size / quant_size:.2f}x")

# 7. 评估量化模型精度
interpreter = tf.lite.Interpreter(model_path=quant_model_path)
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']

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 = input_data.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) == y_test[i]:
        correct += 1

quant_acc = correct / len(x_test)
print(f"Int8量化模型精度: {quant_acc:.4f}")
print(f"精度变化: {quant_acc - float_acc:+.4f}")

运行方式

pip install tensorflow numpy
python mnist_quantization.py

验证方法与判定标准

运行脚本后，观察终端输出。满足以下条件即判定案例通过：

Float32模型精度应高于0.95（训练2轮即可达到）；
量化后模型文件大小明显小于原始模型（通常压缩比在3～5倍）；
量化后模型精度与原始模型相比下降不超过2个百分点（即精度差在-0.02以内）；
程序无异常退出，所有输出数值均为实际运行结果。

关键说明

量化过程中representative_dataset必须提供真实训练数据样本（代码中取前100张），用于校准激活值的动态范围。若数据分布与真实场景差异过大，会导致量化精度显著下降。推理时需手动完成输入量化和输出反量化，这是TFLite int8模型的固定调用方式。

上图展示了本案例的完整处理链路：从Float32模型训练、精度评估，到TFLite int8量化，最终对比两个版本在模型大小和精度上的差异。量化过程通过representative_dataset采样真实数据完成激活值校准，这是保证量化后精度不显著下降的关键步骤。