1. 环境搭建与模型转换入门

任务目标

本案例完成TinyML最小开发环境搭建,并将一个Keras训练的手写数字识别模型转换为TensorFlow Lite格式,验证转换后模型推理结果与原始模型一致。整个过程在PC端完成,不涉及嵌入式硬件部署。

环境依赖

  • Python 3.9+
  • TensorFlow 2.13+(包含TFLite转换器)
  • NumPy

安装命令:pip install tensorflow numpy

完整实现代码

import numpy as np
import tensorflow as tf
from tensorflow import keras

# 1. 训练一个简单模型(MNIST子集,仅用前2000个样本加速)
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train, x_test = x_train[:2000] / 255.0, x_test[:200] / 255.0
y_train, y_test = y_train[:2000], y_test[:200]

model = keras.Sequential([
    keras.layers.Flatten(input_shape=(28, 28)),
    keras.layers.Dense(32, activation='relu'),
    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=64, verbose=0)

# 2. 转换为TFLite格式
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

# 3. 保存转换后的模型
with open('mnist_model.tflite', 'wb') as f:
    f.write(tflite_model)

# 4. 加载TFLite模型并推理
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 取测试集第一个样本验证
test_input = x_test[0:1].astype(np.float32)
interpreter.set_tensor(input_details[0]['index'], test_input)
interpreter.invoke()
tflite_output = interpreter.get_tensor(output_details[0]['index'])

# 5. 对比原始模型与TFLite模型输出
original_output = model.predict(test_input, verbose=0)
print("原始模型预测类别:", np.argmax(original_output[0]))
print("TFLite预测类别:", np.argmax(tflite_output[0]))
print("输出最大差异:", np.max(np.abs(original_output - tflite_output)))

运行方式

将上述代码保存为 tflite_convert_demo.py,在终端执行:

python tflite_convert_demo.py

验证方法与通过标准

  • 程序正常结束且无报错,当前目录生成 mnist_model.tflite 文件(大小约数十KB);
  • 控制台输出的原始模型与TFLite模型预测类别一致(均为0~9中的同一数字);
  • 输出最大差异为远小于1的浮点数(通常<1e-4),说明转换精度损失可忽略。

满足以上三点即判定环境搭建与模型转换流程验证通过。

常见问题

  • 下载MNIST数据集失败:首次运行需联网下载,若网络受限,可手动下载数据集文件并放置于 ~/.keras/datasets/ 目录。
  • 转换警告:若出现TF Lite转换器关于某些算子不支持的警告,本案例所用Dense+Softmax结构不受影响,可忽略。