onnx转TFLite实现本地化部署

将 ONNX 模型转换为 TFLite 模型需要 ​​两个步骤​​：先将 ONNX 转换为 TensorFlow 格式（如 SavedModel），再通过 TensorFlow Lite 转换工具生成 TFLite 模型。以下是详细操作流程（在CPU下即可完成）：

一. 安装依赖库

pip install onnx tf2onnx tensorflow

二. 将 ONNX 转换为 TensorFlow 格式

参考代码：

import tensorflow as tf
from onnx_tf.backend import prepare
import onnx

# 加载 ONNX 模型
onnx_model = onnx.load("input_model.onnx")

# 转换为 TensorFlow SavedModel
tf_rep = prepare(onnx_model)
tf_rep.export_graph("saved_model")  # 输出 SavedModel 格式

三. 将 TensorFlow 模型转为 TFLite

import tensorflow as tf

# 加载 SavedModel
converter = tf.lite.TFLiteConverter.from_saved_model("saved_model")

# 可选：启用量化（减小模型体积）
converter.optimizations = [tf.lite.Optimize.DEFAULT]

# 转换模型
tflite_model = converter.convert()

# 保存 TFLite 模型
with open("output_model.tflite", "wb") as f:
    f.write(tflite_model)

四. 验证准换结果

import tensorflow as tf

# 加载 TFLite 模型
interpreter = tf.lite.Interpreter(model_path="output_model.tflite")
interpreter.allocate_tensors()

# 获取输入/输出张量信息
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 执行推理（示例输入）
input_data = tf.random.normal(input_details[0]['shape'])
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()
output_data = interpreter.get_tensor(output_details[0]['index'])

print("推理结果：", output_data)

五. 将上述步骤整合，完整代码

import tensorflow as tf
from onnx_tf.backend import prepare
import onnx


def onnx2savemodel(onnx_path, save_dir):
    # 加载 ONNX 模型
    onnx_model = onnx.load(onnx_path)

    # 转换为 TensorFlow SavedModel
    tf_rep = prepare(onnx_model)
    tf_rep.export_graph(save_dir)  # 输出 SavedModel 格式


def convert(model_path="./input_model.onnx", save_dir="saved_model"):
    onnx2savemodel(onnx_path, save_dir)

    # 加载 SavedModel 并转换
    converter = tf.lite.TFLiteConverter.from_saved_model("saved_model")
    tflite_model = converter.convert()

    # 保存 TFLite 模型
    with open("output_model.tflite", "wb") as f:
        f.write(tflite_model)


def inference(model_path="output_model.tflite"):
    # 加载 TFLite 模型
    interpreter = tf.lite.Interpreter(model_path=model_path)
    interpreter.allocate_tensors()

    # 获取输入/输出张量信息
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    # 执行推理（示例输入）
    input_data = tf.random.normal(input_details[0]['shape'])
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()
    output_data = interpreter.get_tensor(output_details[0]['index'])

    print("推理结果：", output_data)


if __name__ == "__main__":
    convert()      # 模型转换
    inference()    # 模型推理

OK，就这么轻松愉快的完成了，欢迎关注微信公众号：AIWorkshopLab