4. TensorFlow Lite量化

任务目标

将训练好的Keras模型（以MNIST手写数字分类为例）转换为TensorFlow Lite格式，应用默认的动态范围量化（post-training dynamic range quantization），并验证转换后模型的推理正确性及体积变化。本案例聚焦于转换流程与量化效果验证，不涉及边缘设备实测部署。

环境依赖

需要TensorFlow 2.x（包含TFLite Converter），建议版本≥2.5。无需额外安装包。若需在PC上模拟边缘设备推理，可使用tflite-runtime（可选）。

完整实现代码

import tensorflow as tf
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. 构建并训练一个简单Keras模型（为演示转换，训练1个epoch即可）
model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(28, 28)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(x_train, y_train, epochs=1, batch_size=128, validation_split=0.1, verbose=1)

# 3. 保存Keras模型（H5格式）
model.save('mnist_keras.h5')

# 4. 转换为TFLite（默认动态范围量化）
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]  # 默认量化
tflite_model = converter.convert()

# 5. 保存TFLite模型
with open('mnist_quantized.tflite', 'wb') as f:
    f.write(tflite_model)

# 6. 加载TFLite模型并验证推理
interpreter = tf.lite.Interpreter(model_path='mnist_quantized.tflite')
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 取测试集前100张图片验证
correct = 0
total = 100
for i in range(total):
    input_data = np.expand_dims(x_test[i], axis=0).astype(np.float32)
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()
    output = interpreter.get_tensor(output_details[0]['index'])
    pred = np.argmax(output)
    if pred == y_test[i]:
        correct += 1

print(f'量化模型在{total}张测试图片上的准确率: {correct/total:.2%}')

# 7. 比较模型文件大小
import os
keras_size = os.path.getsize('mnist_keras.h5') / 1024
tflite_size = os.path.getsize('mnist_quantized.tflite') / 1024
print(f'Keras模型大小: {keras_size:.1f} KB')
print(f'TFLite量化模型大小: {tflite_size:.1f} KB')
print(f'体积缩减: {(1 - tflite_size/keras_size):.1%}')

运行方式

将上述代码保存为tflite_quantize.py，在Python环境中直接执行：

python tflite_quantize.py

脚本会自动完成数据加载、模型训练（1个epoch）、转换、保存和验证全流程。

验证方法与判定标准

• 转换成功标志：脚本无异常退出，且生成mnist_quantized.tflite文件。
• 推理正确性：控制台输出的准确率应接近训练时的验证准确率（通常≥90%）。由于仅训练1个epoch且测试100张图片，准确率在85%~95%之间均属正常。
• 体积缩减：量化后的TFLite文件应明显小于原始Keras H5文件（通常缩减约75%~80%）。

若准确率显著低于预期（如<80%），请检查输入数据预处理是否与训练时一致（归一化到[0,1]），以及是否使用了正确的输入shape（1,28,28）。

关键说明

默认的Optimize.DEFAULT会应用动态范围量化：权重从float32转为int8，激活值在推理时动态量化。这种量化方式无需校准数据集，是最简单的TFLite量化入口。若需进一步压缩或硬件加速，可改用tf.lite.Optimize.DEFAULT配合代表性数据集进行全整数量化，但本案例仅覆盖默认量化路径。