2. TFLite模型推理基础

任务目标

本案例演示如何在Python中使用TensorFlow Lite解释器加载一个已转换的.tflite模型，对输入数据执行推理，并获取输出结果。案例以图像分类为场景，使用MobileNetV2模型（已转换为TFLite格式）识别一张测试图片的类别。

环境与依赖

• Python 3.8+
• TensorFlow 2.x（包含TFLite运行时）
• NumPy
• Pillow（用于图片加载）

安装命令：pip install tensorflow numpy pillow

模型与测试图片准备

本案例使用TensorFlow官方提供的MobileNetV2 TFLite模型（输入224×224×3，ImageNet 1000类）。执行以下命令下载模型和测试图片：

wget https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_1.0_224_quant_and_labels.zip
unzip mobilenet_v1_1.0_224_quant_and_labels.zip

解压后得到mobilenet_v1_1.0_224_quant.tflite模型文件和labels.txt标签文件。测试图片使用任意一张224×224的JPEG图片，命名为test.jpg放在同目录下。

完整实现代码

import numpy as np
import tensorflow as tf
from PIL import Image

# 1. 加载TFLite模型
interpreter = tf.lite.Interpreter(model_path="mobilenet_v1_1.0_224_quant.tflite")
interpreter.allocate_tensors()

# 2. 获取输入/输出张量信息
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
print(f"输入张量: {input_details[0]['shape']}, 类型: {input_details[0]['dtype']}")
print(f"输出张量: {output_details[0]['shape']}, 类型: {output_details[0]['dtype']}")

# 3. 预处理输入图片
img = Image.open("test.jpg").resize((224, 224))
input_data = np.array(img, dtype=np.float32)
# 归一化到[0,1]（模型为量化模型，输入范围[0,1]）
input_data = input_data / 255.0
# 增加batch维度: (1, 224, 224, 3)
input_data = np.expand_dims(input_data, axis=0)

# 4. 设置输入张量并执行推理
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()

# 5. 获取输出结果
output_data = interpreter.get_tensor(output_details[0]['index'])
# 量化模型输出为uint8，值范围[0,255]，对应置信度
scores = output_data[0].astype(np.float32) / 255.0

# 6. 解析Top-5结果
with open("labels.txt") as f:
    labels = [line.strip() for line in f.readlines()]

top5_idx = np.argsort(scores)[::-1][:5]
print("\nTop-5 预测结果:")
for i, idx in enumerate(top5_idx):
    print(f"{i+1}. {labels[idx]}: {scores[idx]:.4f}")

运行方式

python tflite_inference.py

验证方法与判定标准

• 程序正常结束：无报错，打印出输入/输出张量信息。
• 输出合理性：Top-5结果中应包含与图片内容相符的类别（例如测试图片为狗，则"golden retriever"等犬类标签应出现在Top-5中），且最高置信度分数应明显高于其他类别（通常>0.5）。
• 数值范围正确：打印的置信度分数均在[0,1]区间内。

若使用非量化模型（如mobilenet_v1_1.0_224.tflite），输入归一化方式不变，但输出为float32，无需除以255，直接使用output_data[0]作为置信度即可。

常见问题

• 输入尺寸不匹配：确保resize为224×224，且通道顺序为RGB。
• 量化模型输入范围：量化TFLite模型输入通常要求[0,1]浮点或[0,255]整数，需根据input_details[0]['dtype']判断：若为float32则归一化到[0,1]；若为uint8则直接使用原始像素值。
• 标签文件行尾：labels.txt每行一个类别名，顺序与模型输出索引一一对应。