深度学习模型部署框架之tflite

---

0 前言

本文是深度学习模型部署框架的第二篇，更多部署框架可以查看本专栏的其他文章。

1 简介

概念：TensorFlow Lite 是TensorFlow在移动和嵌入式设备上的轻量级推理框架，专门针对资源受限的环境进行了优化，以便在手机、嵌入式设备上运行深度学习模型。
优势：缩减了模型的大小和功耗，适合嵌入式等资源受限的设备。
坑点：tfl 目前主要支持 cnn 相关的算子 ，对 rnn 等其他网络中的算子还没有很好的支持。对自家的tf框架支持比较好，对于其他框架支持不佳。

2 模型准备

2.1 导出模型为tfl格式

• python api（推荐）：tf.lite.TFLiteConverter.from_saved_model()

import tensorflow as tf

# 转换模型
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir) # SavedModel目录的路径
tflite_model = converter.convert()

# 这里可以添加其他操作，比如量化、自定义tf操作支持等

# 保存TF Lite模型
with open('model.tflite', 'wb') as f:
  f.write(tflite_model)

• 命令行工具：tflite_convert

$ tflite_convert --help

`--output_file`. Type: string. Full path of the output file.
`--saved_model_dir`. Type: string. Full path to the SavedModel directory.
`--keras_model_file`. Type: string. Full path to the Keras H5 model file.

• python库，可能不适用于某些特定类型的ONNX模型，参考onnx2tflite

from onnx2tflite import onnx_converter

res = onnx_converter(
        onnx_model_path = "/content/denoiser_model_fixed.onnx",
        need_simplify = True,
        output_path = "./models/",
        target_formats = ['tflite'],
    )

• onnx->tf ，参考onnx_tf

# onnx_tf
import onnx
from onnx_tf.backend import prepare

onnx_model_path = "/content/denoiser_model_fixed.onnx"
onnx_model = onnx.load(onnx_model_path)

tf_rep = prepare(onnx_model)
tf_rep.export_graph("/content/denoiser_model.pb")

• 注意，如果要将其他格式的模型转为tfl，需要other format->onnx->tf->tfl，主要难点在于onnx->tf

2.2 导出时注意事项

• 动态形状的支持，tf转tflite时，由于tf.TensorListReserve操作需要在转换过程中具有静态的element_shape，但原生模型是动态的，所以在TensorFlow Lite 转换器中启用选择性TF操作支持，通过设置converter.target_spec.supported_ops = [ tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS ]实现。这里只是提到一个我实践中遇到的案例，更多的操作可以自定义。
• 环境配置：转换模型时环境的配置很容易出错，目前实践成功的环境配置为tensorflow=2.10.0，tensorflow-probability=0.16.0，protobuf=3.20.3，太低或太高会报错。
• 模型的opset版本：转onnx时，opset_version的设置会影响后续的过程，最终设置opset_version=12可以正常运行
• 当模型比较小的时候，导出为tfl时需要仔细验证模型前后的输出误差，一般来说模型越小，转换以为tfl后输出的误差越大，如果对误差不能接受，可以选择更保守一些的量化方案，如float16量化。