TensorFlow-ONNX: 将TensorFlow模型转换为ONNX格式

tensorflow-onnx

TensorFlow-ONNX简介

TensorFlow-ONNX(tf2onnx)是一个开源项目,旨在将TensorFlow、Keras、TensorFlow.js和TFLite模型转换为ONNX(开放神经网络交换)格式。ONNX是一种用于表示机器学习模型的开放标准格式,支持多种深度学习框架之间的互操作性。通过将TensorFlow模型转换为ONNX格式,可以在不同的深度学习框架和推理引擎中使用这些模型,提高了模型的可移植性和灵活性。

TensorFlow-ONNX支持多种类型的TensorFlow模型,包括:

TensorFlow 1.x和2.x版本的模型
Keras模型
TensorFlow.js模型
TFLite模型

它可以通过命令行接口或Python API进行模型转换。

支持的版本

TensorFlow-ONNX支持以下版本:

ONNX

支持并测试ONNX opset 14到18
opset 6到13应该也可以工作,但未经全面测试
默认使用opset 15生成ONNX图

TensorFlow

支持TensorFlow 1.x和2.x
在TensorFlow 1.15或更高版本上进行测试

Python

支持Python 3.7-3.10

安装

安装TensorFlow-ONNX有以下几种方式:

通过pip安装:

pip install -U tf2onnx

从GitHub安装最新版本:

pip install git+https://github.com/onnx/tensorflow-onnx

从源代码构建安装(用于开发):

git clone https://github.com/onnx/tensorflow-onnx
cd tensorflow-onnx
python setup.py install

使用方法

命令行使用

最基本的用法是通过命令行将保存的TensorFlow模型转换为ONNX:

python -m tf2onnx.convert 
    --saved-model tensorflow-model-path 
    --output model.onnx

主要参数说明:

--saved-model: TensorFlow SavedModel格式模型的路径
--output: 输出的ONNX模型文件路径
--opset: 指定ONNX opset版本,默认为15
--inputs, --outputs: 指定输入和输出节点名称
--tag, --signature_def: 用于SavedModel的tag和signature
--target: 指定目标运行时,如果模型将在Windows ML上运行需要指定

更多参数选项可以参考CLI文档。

Python API使用

TensorFlow-ONNX也提供了Python API用于程序化转换模型。主要的转换函数包括:

from_keras: 转换Keras模型
from_function: 转换tf.function
from_graph_def: 转换TensorFlow GraphDef
from_tflite: 转换TFLite模型

以下是一个使用from_keras转换Keras模型的示例:

import tf2onnx

model_proto, _ = tf2onnx.convert.from_keras(keras_model, 
                                            opset=12,
                                            input_signature=(tf.TensorSpec((None, 224, 224, 3), tf.float32),),
                                            output_path="model.onnx")

更多API使用示例可以参考Python API文档。

工作原理

TensorFlow-ONNX的转换过程主要包括以下步骤:

从冻结的TensorFlow图开始
将TensorFlow protobuf格式转换为ONNX protobuf格式
重写子图以匹配ONNX操作
处理单个操作,映射TensorFlow操作到ONNX操作
优化生成的ONNX图
进行最终的拓扑排序

在这个过程中,TensorFlow-ONNX需要处理一些主要挑战:

将TensorFlow数据类型映射到ONNX等效类型
将TensorFlow中作为输入传递的参数(如shape)转换为ONNX中的属性
识别并替换TensorFlow中由多个简单操作组成的复合操作
处理TensorFlow的NHWC数据格式和ONNX的NCHW格式之间的转换
使用其他ONNX操作来实现ONNX中不支持的操作(如relu6)

通过这些步骤,TensorFlow-ONNX能够将复杂的TensorFlow模型准确地转换为ONNX格式,同时保持模型的功能和性能。

扩展TensorFlow-ONNX

如果你想为TensorFlow-ONNX添加新的转换支持,可以遵循以下步骤:

检查新操作是否适合现有的映射,如果是,只需将其添加到_OPS_MAPPING字典中
如果需要额外处理,创建新的映射函数
对于由多个操作组成的TensorFlow操作,考虑使用图重写
在tests/test_backend.py中添加单元测试
如果有使用新操作的预训练模型,考虑将其添加到test/run_pretrained_models.py中

总结

TensorFlow-ONNX是一个强大而灵活的工具,能够将各种TensorFlow模型转换为ONNX格式。它支持广泛的模型类型和版本,提供了简单的命令行接口和Python API。通过将TensorFlow模型转换为ONNX,可以在不同的深度学习框架和推理引擎中使用这些模型,极大地提高了模型的可移植性和互操作性。

对于想要在不同平台和框架之间共享和部署TensorFlow模型的开发者和研究人员来说,TensorFlow-ONNX是一个非常有价值的工具。它解决了深度学习生态系统中的一个重要问题,促进了不同框架之间的模型交换和协作。

随着深度学习技术的不断发展,TensorFlow-ONNX也在持续更新和改进,以支持最新的TensorFlow特性和ONNX标准。社区的贡献和反馈对项目的发展起着重要作用,欢迎感兴趣的开发者参与到TensorFlow-ONNX的开发和改进中来。

TensorFlow to ONNX conversion process