📚 课程目录第4章 TFLite模型在嵌入式设备上的缺陷分类移植

4. TFLite模型在嵌入式设备上的缺陷分类移植

任务目标

将训练好的工业缺陷分类模型（Keras .h5）转换为TensorFlow Lite格式，并在嵌入式Linux设备（如树莓派4B）上完成推理移植，实现对传送带工件图像的实时缺陷分类。本案例聚焦于转换、量化、部署与验证的完整闭环。

环境与依赖

PC端：TensorFlow 2.8+（含TFLite Converter）、Python 3.8+、numpy、Pillow
嵌入式端：Python 3.7+、tflite-runtime（或完整TensorFlow）、numpy、Pillow
模型来源：自训练的Keras模型（输入224×224×3，输出6类缺陷）

关键完整代码

# ============ 文件1: convert_tflite.py（PC端执行） ============
import tensorflow as tf
import numpy as np

# 1. 加载Keras模型并转换（含量化校准数据）
model = tf.keras.models.load_model('defect_model.h5')

# 准备100张代表性校准图像（灰度转RGB，归一化到[0,1]）
def representative_dataset():
    for i in range(100):
        img = tf.keras.preprocessing.image.load_img(
            f'calib_imgs/img_{i}.jpg', target_size=(224, 224))
        arr = tf.keras.preprocessing.image.img_to_array(img) / 255.0
        yield [arr[np.newaxis, ...].astype(np.float32)]

# 2. 转换：使用INT8全整数量化（嵌入式端无浮点加速时推荐）
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_dataset
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8
tflite_model = converter.convert()

# 3. 保存转换后的模型
with open('defect_model_int8.tflite', 'wb') as f:
    f.write(tflite_model)
print(f'转换完成，模型大小: {len(tflite_model)/1024:.1f} KB')

# ============ 文件2: deploy_inference.py（嵌入式设备执行） ============
import numpy as np
from PIL import Image
try:
    from tflite_runtime.interpreter import Interpreter
except ImportError:
    from tensorflow.lite.python.interpreter import Interpreter

# 类别标签（与训练时一致）
CLASS_NAMES = ['scratch', 'dent', 'crack', 'stain', 'burr', 'normal']

class DefectClassifier:
    def __init__(self, model_path):
        self.interpreter = Interpreter(model_path=model_path)
        self.interpreter.allocate_tensors()
        self.input_details = self.interpreter.get_input_details()
        self.output_details = self.interpreter.get_output_details()
        # 获取量化参数
        self.input_scale, self.input_zero = self.input_details[0]['quantization']
        self.output_scale, self.output_zero = self.output_details[0]['quantization']

    def preprocess(self, image_path):
        """读取图像并转换为uint8输入格式"""
        img = Image.open(image_path).resize((224, 224)).convert('RGB')
        # 反量化：将[0,1]浮点映射到uint8范围
        arr = np.asarray(img, dtype=np.float32) / 255.0
        arr = arr / self.input_scale + self.input_zero
        return arr.astype(np.uint8)[np.newaxis, ...]

    def predict(self, image_path):
        """执行推理并返回类别索引与置信度"""
        input_data = self.preprocess(image_path)
        self.interpreter.set_tensor(self.input_details[0]['index'], input_data)
        self.interpreter.invoke()
        output = self.interpreter.get_tensor(self.output_details[0]['index'])
        # 输出反量化到浮点概率
        probs = (output.astype(np.float32) - self.output_zero) * self.output_scale
        probs = np.exp(probs - np.max(probs))  # softmax稳定化
        probs = probs / np.sum(probs)
        return int(np.argmax(probs)), float(np.max(probs))

if __name__ == '__main__':
    clf = DefectClassifier('defect_model_int8.tflite')
    idx, conf = clf.predict('test_sample.jpg')
    print(f'预测类别: {CLASS_NAMES[idx]}, 置信度: {conf:.3f}')

运行与部署步骤

PC端转换：将训练好的defect_model.h5与100张校准图放入calib_imgs/目录，执行python convert_tflite.py，生成defect_model_int8.tflite。
拷贝到设备：将.tflite文件与deploy_inference.py通过scp或U盘拷贝至嵌入式设备。
安装运行时：在设备上执行pip install tflite-runtime（若使用完整TF则无需此步）。
执行推理：将一张测试图像命名为test_sample.jpg放在同目录，运行python deploy_inference.py。

验证方法与通过标准

转换验证：PC端输出模型大小（INT8量化后通常为原模型1/4左右），且无报错。
推理一致性：在PC端用原始Keras模型对同一张test_sample.jpg预测，记录类别；再在嵌入式端运行本代码，两者预测类别应一致（置信度允许±0.05差异）。
性能指标：在树莓派4B上单次推理耗时应在200ms以内（可通过在predict前后加time.time()测量），满足产线节拍要求。
边界测试：分别用6类缺陷样本各测试5张，分类准确率应≥95%（与原始模型在验证集上的准确率差异≤2%）。

常见问题

量化后精度下降：增加校准集数量至200张以上，并确保校准图覆盖各类缺陷的灰度分布范围。
设备端无tflite-runtime：若设备为ARM64架构，可下载预编译wheel包安装；或直接安装完整版TensorFlow。
输入尺寸不匹配：确认resize尺寸与训练时一致（本案例为224×224），否则需修改预处理代码。