onnx转tflite

Index

方法一：onnx-tf

参考：https://zhuanlan.zhihu.com/p/363317178

环境配置

（使用anaconda配置更方便些）
onnx-tf这个库目前没有主动维护，现在的版本是1.10.0，这里就使用1.10.0版本：
关于环境配置要求，本文都贴出了参考链接，若要使用其他版本，请参考相应链接检查自己版本。
当然，各个包肯定能兼容其他版本的，我这里大都是按照最低兼容版本配置的。
onnx-tf==1.10.0 搭配 tensorflow==2.8.0, onnx==1.10.2
https://github.com/onnx/onnx-tensorflow
https://github.com/onnx/onnx-tensorflow/blob/main/ONNX_VERSION_NUMBER
与tensorflow==2.8.0 需要搭配 tensorflow-probability==0.16.0 和JAX 0.3.0 (与GPU相关)；
https://github.com/tensorflow/probability/releases
支持需要python3.6-3.9，本文测试python环境为3.8。
https://www.tensorflow.org/install/pip?hl=zh-cn
同时与2.8tensorflow兼容的tensorflow-addons版本如下，这里选择0.18.0
https://github.com/tensorflow/addons

onnx==1.10.2和onnxruntime兼容的版本有1.9和1.10
https://onnxruntime.ai/docs/reference/compatibility.html

测试过程如果遇到protobuf或者numpy的报错，可以降一下版本，本人最终配置成功后的环境如下：

keras==2.8.0  # 装tensorflow就会有，可以不用管

numpy==1.21.0

onnx==1.10.2
onnx-tf==1.10.0
onnxruntime==1.9.0

protobuf==3.19.0

tensorflow==2.8.0
tensorflow-addons==0.18.0
tensorflow-probability==0.16.0

onnx文件

导出onnx文件注意下opset version，需要设置为与onnx兼容的版本，理论上是可以向下兼容至7。
其他版本的参考上面的给出的地址：https://onnxruntime.ai/docs/reference/compatibility.html

本文用paddle2onnx最新版本(1.0.5)转换的onnx文件，发现onnx-tf对Squeeze算子支持会报version相关的错，在导出onnx文件时设置opset_version为11能转换成功(15-13都会报错，12没试）。
对于不知道opset_version的onnx文件，load之后可以查看：

onnx_model = onnx.load(ONNX_PATH)
print(onnx_model.opset_import)

另外可以将onnx模型输入固定为静态shape，若为动态shape输入，在netron中看到相应的维度为-1，这以paddle2onnx工具为例，修改输入shape：

python -m paddle2onnx.optimize --input_model yolo5n.onnx \
                               --output_model new_yolo5n.onnx \
                               --input_shape_dict "{'image':[1,3,640,640]}"

检查onnx是否有错误：

ONNX_PATH = 'onnx文件路径'
onnx.checker.check_model(ONNX_PATH)

onnx去冗余（可选）pip install onnx-simplifier

import onnx
from onnxsim import simplify

# load your predefined ONNX model
model = onnx.load(path + model_name + '.onnx')

# convert model
model_simp, check = simplify(model)

assert check, "Simplified ONNX model could not be validated"

# use model_simp as a standard ONNX model object

转换代码：

终端调用转换的命令可到github项目页查看，这里给出python代码版本：

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

TF_PATH = "weights/test_out/yolov5.pb"  # 转tensorflow的pb文件的存储路径
ONNX_PATH = "weights/test/yolov5_best.onnx"  # 需要转换的onnx文件路径
TFLITE_PATH = "weights/test_out/yolov5.tflite"  # 输出的tflite文件路径

"""onnx=>.pb文件"""
onnx_model = onnx.load(ONNX_PATH)  # load onnx model
tf_rep = prepare(onnx_model)  # creating TensorflowRep object
tf_rep.export_graph(TF_PATH)

""".pb=>tflite"""
converter = tf.lite.TFLiteConverter.from_saved_model(TF_PATH)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tf_lite_model = converter.convert()
with open(TFLITE_PATH, 'wb') as f:
    f.write(tf_lite_model)

最后在控制台有输出模型信息则模型转换成功：
Estimated count of arithmetic ops: xxxx ops, equivalently xxxx MACs

方法二：onnx2tflite

项目地址：https://github.com/MPolaris/onnx2tflite

环境要求：

onnx
onnxruntime
onnx-simplifier
numpy
tensorflow>=2.5
opencv-python

这里没有对环境配置做过多探究，因为我尝试方法一的时候配置好了环境，还没有配置环境的话，可以试试克隆项目后pip install -r requirements.txt。

转换代码：

使用方法就比较简单了，这里贴出python代码，详情查看上面项目地址。
实际测试torch转换的onnx文件没有问题，这里opset version也是11。

import sys
sys.path.append("onnx2tflite")  # onnx2tflite的地址

from converter import onnx_converter
onnx_path = "weights/test/yolov5_best.onnx"  # 需要转换的onnx文件位置
onnx_converter(
    onnx_model_path = onnx_path,
    need_simplify = True,
    output_path = "weights/test_out/",  # 输出的tflite存储路径
    target_formats = ['tflite'], # or ['keras'], ['keras', 'tflite']
    weight_quant = False,
    int8_model = False,
    int8_mean = None,
    int8_std = None,
    image_root = None
)

对于paddle2onnx转换的onnx文件则有些许不同：
由于paddle2onnx处理后的onnx没有initliazer，所有的权重转成了constant节点，不能直接使用onnx2tflite，可以使用onnx-simplifier处理后就会转成initliazer，就能使用onnx2tflite了。
https://github.com/PaddlePaddle/Paddle2ONNX/issues/741

import sys
sys.path.append("onnx2tflite")
from converter import onnx_converter
onnx_path = "weights/test/new_yolo5n.onnx"  # onnx地址


# pip install onnx-simplifier
import onnx
from onnxsim import simplify
model = onnx.load(onnx_path)
model_simp, check = simplify(model)
assert check, "Simplified ONNX model could not be validated"

simp_onnx_path = 'weights/test/simp_yolo5n.onnx'  # simp后的onnx存放地址
onnx.save(model_simp, simp_onnx_path)  # 保存文件

onnx_converter(
    onnx_model_path = simp_onnx_path,
    need_simplify = True,
    output_path = "weights/test_out/",
    target_formats = ['tflite'],
    weight_quant = False,
    int8_model = False,
    int8_mean = None,
    int8_std = None,
    image_root = None
)

两种方法的区别

onnx-tf

由于pytorch的输入是NCHW，转成ONNX也是NCHW，再使用onnx-tf转成tflite时，输入也是NCHW，所以在某些需要以NHWC为输入的算子上（如conv），就会在该算子的前后分别多出一个transpose算子（第一个用于NCHW->NHWC，第二个用于NHWC->NCHW），这也是onnx-tf转换的生硬之处，多出的算子会对推理速度有一些影响。

onnx2tflite，则能够将NCWH格式的通道变换为tensorflow的NWHC格式，自然也就省去了conv两头的transpose算子。

当然后处理时，注意两种方式的输入输出shape区别：

Index

方法一：onnx-tf

环境配置

onnx文件

转换代码：

方法二：onnx2tflite

环境要求：

转换代码：

两种方法的区别

猜你喜欢