tf2onnx - 将 TensorFlow, Keras, Tensorflow.js 和 Tflite 模型转换为 ONNX.
tf2onnx 通过命令行或 python api 将 TensorFlow (tf-1.x 或 tf-2.x), keras, tensorflow.js 和 tflite 模型转换为 ONNX。
注:tensorflow.js 支持刚刚添加。我们使用来自 tfhub 的许多 tfjs 模型进行了测试,但应视为实验性功能。
TensorFlow 拥有比 ONNX 更多的操作,有时将模型映射到 ONNX 会产生问题。
您可以在此处找到支持的 TensorFlow 操作及其映射到 ONNX 的列表。
我们尝试在此处故障排除指南记录常见问题。
构建类型 | 操作系统 | Python | TensorFlow | ONNX 操作集 | 状态 |
---|---|---|---|---|---|
单元测试 - 基本 | Linux, Windows | 3.7-3.10 | 1.15, 2.9-2.15 | 14-18 | |
单元测试 - 完整 | Linux, Windows | 3.7-3.10 | 1.15, 2.9-2.15 | 14-18 |
支持的版本
ONNX
tf2onnx 将使用您系统上安装的 ONNX 版本,并在未找到时安装最新的 ONNX 版本。
我们支持并测试 ONNX 操作集-14 到 操作集-18。操作集-6 到 操作集-13 应该可以工作,但我们不测试它们。
默认情况下,我们使用 操作集-15
生成 ONNX 图。
如果希望使用特定的操作集生成图,请在命令行中使用 --opset
,例如--opset 15
。
TensorFlow
我们支持 tf-1.x 图
和 tf-2.x
。为了保持我们的测试矩阵可管理,我们在 tf-1.15 或更高版本
上测试 tf2onnx。
在 tf-2.x 版本下运行时,tf2onnx 将使用 tensorflow V2 控制流。
您可以在 tf-1.x 或 tf-2.x 之上安装 tf2onnx。
Python
我们支持 Python 3.7-3.10
。
先决条件
安装 TensorFlow
如果您尚未安装 TensorFlow,请安装所需的 TensorFlow 构建,例如:
pip install tensorflow
(可选)安装运行时
如果希望运行测试,请安装一个可以运行 ONNX 模型的运行时。例如:
ONNX 运行时(适用于 Linux, Windows 和 Mac):
pip install onnxruntime
安装
从 pypi 安装
pip install -U tf2onnx
从 GitHub 安装最新版本
pip install git+https://github.com/onnx/tensorflow-onnx
从源代码构建并安装最新版本(用于开发)
git clone https://github.com/onnx/tensorflow-onnx
在 tensorflow-onnx 文件夹中安装依赖项后,调用:
python setup.py install
或
python setup.py develop
tensorflow-onnx 需要 onnx-1.9 或更高版本,并在需要时安装/升级 onnx。
要创建用于分发的 wheel 包:
python setup.py bdist_wheel
入门指南
要开始使用 tensorflow-onnx
,运行 t2onnx.convert
命令,提供:
- 您的 TensorFlow 模型路径(模型格式为
saved model
) - ONNX 输出文件名:
python -m tf2onnx.convert --saved-model tensorflow-model-path --output model.onnx
上述命令使用默认的 15
作为 ONNX 操作集。如果需要更新的操作集,或希望限制模型使用较旧的操作集,则可以在命令中提供 --opset
参数。如果不确定使用哪个操作集,请参考ONNX 操作符文档。
python -m tf2onnx.convert --saved-model tensorflow-model-path --opset 18 --output model.onnx
如果您的 TensorFlow 模型格式不是 saved model
,则需要提供模型图的输入和输出:
对于 checkpoint
格式:
python -m tf2onnx.convert --checkpoint tensorflow-model-meta-file-path --output model.onnx --inputs input0:0,input1:0 --outputs output0:0
对于 graphdef
格式:
python -m tf2onnx.convert --graphdef tensorflow-model-graphdef-file --output model.onnx --inputs input0:0,input1:0 --outputs output0:0
如果您的模型是 checkpoint
或 graphdef
格式,且您不知道模型的输入和输出节点,可以使用 summarize_graph TensorFlow 工具。summarize_graph
工具需要从源码下载并构建。如果可以从模型提供者处获得 saved model
格式的模型,我们建议这样做。
您可以在此处找到 ssd-mobilenet 的端到端教程 here
我们最近添加了对 tflite 的支持。您可以通过命令行将tflite
模型转换,例如:
python -m tf2onnx.convert --opset 16 --tflite tflite--file --output model.onnx
CLI 参考
python -m tf2onnx.convert
--saved-model SOURCE_SAVED_MODEL_PATH |
--checkpoint SOURCE_CHECKPOINT_METAFILE_PATH |
--tflite TFLITE_MODEL_PATH |
--tfjs TFJS_MODEL_PATH |
--input | --graphdef SOURCE_GRAPHDEF_PB
--output TARGET_ONNX_MODEL
[--inputs GRAPH_INPUTS]
[--outputs GRAPH_OUTPUS]
[--inputs-as-nchw inputs_provided_as_nchw]
[--outputs-as-nchw outputs_provided_as_nchw]
[--opset OPSET]
[--dequantize]
[--tag TAG]
[--signature_def SIGNATURE_DEF]
[--concrete_function CONCRETE_FUNCTION]
[--target TARGET]
[--extra_opset list-of-extra-opset]
[--custom-ops list-of-custom-ops]
[--load_op_libraries tensorflow_library_path]
[--large_model]
[--continue_on_error]
[--verbose]
[--output_frozen_graph]
参数
--saved-model
以 saved_model 格式的 TensorFlow 模型。我们期望 saved_model 目录的路径。
--checkpoint
以 checkpoint 格式的 TensorFlow 模型。我们期望 .meta 文件的路径。
--input 或 --graphdef
以 graphdef 文件形式的 TensorFlow 模型。
--tfjs
通过提供 .tfjs 文件路径转换 tensorflow.js 模型。不需要指定输入/输出。
--tflite
通过提供 .tflite 文件路径转换 tflite 模型。不需要指定输入/输出。
--output
目标 onnx 文件路径。
--inputs、--outputs
TensorFlow 模型的输入/输出名称,可以通过 summarize graph tool 找到。这些名称通常以 :0
结尾,例如 --inputs input0:0,input1:0
。对于 saved-model 格式的模型,不需要输入和输出。某些模型指定了无法映射到 onnx 的未知阶和尺寸占位符。在这些情况下,可以在输入名称后添加形状,使用一对 []
括起,例如 --inputs X:0[1,28,28,3]
。使用 -1 表示未知维度。
--inputs-as-nchw
默认情况下,我们保留 TensorFlow 模型给定的输入图像格式(nchw
或 nhwc
)。如果您的主机(例如 Windows)原生格式为 nchw 而模型是为 nhwc 编写,那么--inputs-as-nchw
tensorflow-onnx 将转置输入。这样做对于应用程序很方便,转换器在许多情况下可以优化删除转置。例如 --inputs input0:0,input1:0 --inputs-as-nchw input0:0
假设图像以 nchw 格式传递到 input0:0
,而给定的 TensorFlow 模型使用 nhwc。
--outputs-as-nchw
类似于 --inputs-as-nchw
的使用。默认情况下,我们会保留输出的格式(nchw
或 nhwc
)与 TensorFlow 模型显示的一致。如果你的主机原生格式是 nchw
而模型是 nhwc
,使用 --outputs-as-nchw
tensorflow-onnx 会转置输出并优化掉转置。例如,--outputs output0:0,output1:0 --outputs-as-nchw output0:0
会将 output0:0
更改为 nchw
,而给定的 TensorFlow 模型使用 nhwc
。
--ignore_default, --use_default
ONNX 需要图输入的默认值是常量,而 Tensorflow 的 PlaceholderWithDefault 操作符接受计算默认值。要转换此类模型,传递以逗号分隔的节点名称列表给 ignore_default 和/或 use_default 标志。匹配名称的 PlaceholderWithDefault 节点将分别被替换为 Placeholder 或 Identity 操作符。
--opset
默认情况下我们使用 opset 15 生成图。通过指定 --opset
,用户可以覆盖默认值以生成具有所需 opset 的图。例如 --opset 17
将创建一个只使用 opset 17 中可用操作的 onnx 图。由于较旧的 opsets 在大多数情况下拥有较少的操作,一些模型可能无法在较旧的 opset 上进行转换。
--dequantize
(这是实验性的,仅支持 tflite)
从量化的 tflite 模型生成 float32 模型。从量化边界检测 ReLU 和 ReLU6 操作。
--tag
仅在参数 --saved_model
有效。指定要使用的 saved_model 中的标签。典型值是 'serve'。
--signature_def
仅在参数 --saved_model
有效。指定在指定的 --tag 值内使用的签名。典型值是 'serving_default'。
--concrete_function
(这是实验性的,仅适用于 TF2.x 模型)
仅在参数 --saved_model
有效。如果一个模型包含一个具体函数列表,在具体函数名为 __call__
下(可以使用命令 saved_model_cli show --all
查看),此参数是一个从0开始的整数,指定列表中的哪个函数应转换。此参数优先于 --signature_def
,后者将被忽略。
--target
一些模型需要特殊处理才能在某些运行时上运行。特别地,模型可能使用不支持的数据类型。通过 --target TARGET
激活解决方法。目前支持的值列在这个 wiki 中。如果你的模型将在 Windows ML 上运行,你应该指定适当的目标值。
--extra_opset
如果你想使用现有的自定义操作转换 TF 模型,这可以指定相应的域和版本。格式是域和版本的逗号分隔映射,例如:ai.onnx.contrib:1
。
--custom-ops
如果模型包含 onnx 运行时无法识别的操作,你可以为这些操作附加自定义操作域,以便运行时仍然可以打开模型。格式是 tf 操作名称到域的逗号分隔映射格式 OpName:domain。如果只提供操作名称(没有冒号),则使用默认域 ai.onnx.converters.tensorflow
。
--load_op_libraries
在转换之前加载逗号分隔的 tensorflow 插件/操作库列表。
--large_model
(仅适用于 TF2.x 模型)
仅在参数 --saved_model
有效。当设置时,创建一个 zip 文件,包含 ONNX protobuf 模型和存储在外部的大张量值。这允许转换大小超过 2 GB 的模型。
--continue_on_error
在错误时继续运行转换,忽略图循环,以便报告所有缺少的操作和错误。
--verbose
调试目的的详细输出。
--output_frozen_graph
将冻结的和优化的 tensorflow 图保存到一个文件以进行调试。
获取图输入和输出的工具
要查找 TensorFlow 图的输入和输出,模型开发者将知道或可以参考 TensorFlow 的 summarize_graph 工具,例如:
summarize_graph --in_graph=tests/models/fc-layers/frozen.pb
测试
测试有两种类型。
单元测试
python setup.py test
验证预训练的 TensorFlow 模型
python tests/run_pretrained_models.py
usage: run_pretrained_models.py [-h] [--cache CACHE] [--tests TESTS] [--backend BACKEND] [--verbose] [--debug] [--config yaml-config]
optional arguments:
-h, --help show this help message and exit
--cache CACHE pre-trained models cache dir
--tests TESTS tests to run
--backend BACKEND backend to use
--config yaml config file
--verbose verbose output, option is additive
--opset OPSET target opset to use
--perf csv-file capture performance numbers for tensorflow and onnx runtime
--debug dump generated graph with shape info
run_pretrained_models.py
将运行 TensorFlow 模型,捕获 TensorFlow 输出,并在转换模型后对指定的 ONNX 后端运行相同的测试。
如果指定了 --perf csv-file
选项,我们将捕获 tensorflow 和 onnx 运行时的推理时间,并将结果写入给定的 csv 文件。
你可以这样调用它:
python tests/run_pretrained_models.py --backend onnxruntime --config tests/run_pretrained_models.yaml --perf perf.csv
保存预训练模型的工具
我们提供了一个 实用程序 来保存预训练模型及其配置。在你的最后一个测试周期中,放置 save_pretrained_model(sess, outputs, feed_inputs, save_dir, model_name)
,预训练模型和配置将被保存在 save_dir/to_onnx
下。更多信息请参考 tools/save_pretrained_model.py 中的示例。注意,Tensorflow 的最低要求版本是 r1.6。
Python API 参考
在 tf2onnx-1.8.4 中我们更新了 API。我们的旧 API 仍然有效——你可以在这里找到文档。
from_keras(tf-2.0 及更新)
import tf2onnx
model_proto, external_tensor_storage = tf2onnx.convert.from_keras(model,
input_signature=None, opset=None, custom_ops=None,
custom_op_handlers=None, custom_rewriter=None,
inputs_as_nchw=None, outputs_as_nchw=None, extra_opset=None,
shape_override=None, target=None, large_model=False, output_path=None)
Args:
model: 我们想转换的 tf.keras 模型
input_signature: 一个 tf.TensorSpec 或一个定义输入形状/数据类型的 numpy 数组
opset: 用于 ONNX 模型的 opset,默认是最新的
custom_ops: 如果一个模型包含 onnx 运行时无法识别的操作,你可以为这些操作附加自定义操作域,以便运行时仍然可以打开模型。类型是一个字典 `{op 名称: 域}`.
target: 应用于帮助某些平台的解决方法列表
custom_op_handlers: 自定义操作处理程序的字典
custom_rewriter: 自定义图重写器列表
extra_opset: 额外的 opset 列表,例如自定义操作使用的 opset
shape_override: 覆盖 tensorflow 提供的形状的输入字典
inputs_as_nchw: 将列表中的输入从 nhwc 转置为 nchw
outputs_as_nchw: 将列表中的输出从 nhwc 转置为 nchw
large_model: 使用 ONNX 外部张量存储格式
output_path: 保存模型到 output_path
Returns:
一个 ONNX model_proto 和一个 external_tensor_storage 字典。
请参阅 tutorials/keras-resnet50.ipynb 了解一个端到端示例。
from_function(tf-2.0 及更新)
import tf2onnx
model_proto, external_tensor_storage = tf2onnx.convert.from_function(function,
input_signature=None, opset=None, custom_ops=None,
custom_op_handlers=None, custom_rewriter=None, inputs_as_nchw=None,
outputs_as_nchw=None, extra_opset=None, shape_override=None,
target=None, large_model=False, output_path=None)
参数:
function:我们想要转换的 tf.function
input_signature:一个定义输入的形状/数据类型的 tf.TensorSpec 或 numpy 数组
opset:要用于 ONNX 模型的 opset ,默认是最新的
custom_ops:如果模型包含 onnx runtime 无法识别的操作,
你可以用自定义操作域标记这些操作,以便 runtime 仍然可以打开模型。类型是一个字典 {操作名: 域}
。
target:应用于帮助特定平台的解决方法列表
custom_op_handlers:自定义操作处理程序的字典
custom_rewriter:自定义图形重写列表
extra_opset:额外的 opset 列表,例如自定义操作使用的 opset
shape_override:覆盖 tensorflow 给定形状的输入字典
inputs_as_nchw:将列表中的输入从 nhwc 转置为 nchw
outputs_as_nchw:将列表中的输出从 nhwc 转置为 nchw
large_model:使用 ONNX 外部张量存储格式
output_path:将模型保存到 output_path
返回: 一个 ONNX model_proto 和一个 external_tensor_storage 字典。
from_graph_def
import tf2onnx
model_proto, external_tensor_storage = tf2onnx.convert.from_graph_def(graph_def,
name=None, input_names=None, output_names=None, opset=None,
custom_ops=None, custom_op_handlers=None, custom_rewriter=None,
inputs_as_nchw=None, outputs_as_nchw=None, extra_opset=None,
shape_override=None, target=None, large_model=False,
output_path=None)
参数:
graph_def:我们想要转换的 graph_def
input_names:输入名的列表
output_names:输出名的列表
name:图的名称
opset:要用于 ONNX 模型的 opset,默认为最新的
target:应用于帮助特定平台的解决方法列表
custom_op_handlers:自定义操作处理程序的字典
custom_rewriter:自定义图形重写列表
extra_opset:额外的 opset 列表,例如自定义操作使用的 opset
shape_override:覆盖 tensorflow 给定形状的输入字典
inputs_as_nchw:将列表中的输入从 nhwc 转置为 nchw
outputs_as_nchw:将列表中的输出从 nhwc 转置为 nchw
large_model:使用 ONNX 外部张量存储格式
output_path:将模型保存到 output_path
返回:
一个 ONNX model_proto 和一个 external_tensor_storage 字典。
### from_tflite
import tf2onnx
model_proto, external_tensor_storage = tf2onnx.convert.from_tflite(tflite_path, input_names=None, output_names=None, opset=None, custom_ops=None, custom_op_handlers=None, custom_rewriter=None, inputs_as_nchw=None, outputs_as_nchw=None, extra_opset=None, shape_override=None, target=None, large_model=False, output_path=None):
参数:
tflite_path:tflite 模型文件的完整路径
input_names:输入名的列表
output_names:输出名的列表
opset:要用于 ONNX 模型的 opset,默认为最新的
custom_ops:如果模型包含 onnx runtime 无法识别的操作,
你可以用自定义操作域标记这些操作,以便 runtime 仍然可以打开模型。类型是一个字典 {操作名: 域}
。
custom_op_handlers:自定义操作处理程序的字典
custom_rewriter:自定义图形重写列表
inputs_as_nchw:将列表中的输入从 nhwc 转置为 nchw
outputs_as_nchw:将列表中的输出从 nhwc 转置为 nchw
extra_opset:额外的 opset 列表,例如自定义操作使用的 opset
shape_override:覆盖 tensorflow 给定形状的输入字典
target:应用于帮助特定平台的解决方法列表
large_model:使用 ONNX 外部张量存储格式
output_path:将模型保存到 output_path
返回: 一个 ONNX model_proto 和一个 external_tensor_storage 字典。
创建自定义操作映射从 python
对于需要图形重写或输入/属性重写的复杂自定义操作,使用 python 接口插入自定义操作将是完成任务的最简单方法。 一个名为 name->custom_op_handler 的字典可以传递给 tf2onnx.tfonnx.process_tf_graph。如果在图中找到操作名,处理程序将可以访问所有内部结构并进行所需的重写。例如 examples/custom_op_via_python.py:
import tensorflow as tf
import tf2onnx
from onnx import helper
_TENSORFLOW_DOMAIN = "ai.onnx.converters.tensorflow"
def print_handler(ctx, node, name, args):
# 使用 Identity 替换 tf.Print()
# T output = Print(T input, data, @list(type) U, @string message, @int first_n, @int summarize)
# 变为:
# T output = Identity(T Input)
node.domain = _TENSORFLOW_DOMAIN
del node.input[1:]
return node
with tf.Session() as sess:
x = tf.placeholder(tf.float32, [2, 3], name="input")
x_ = tf.add(x, x)
x_ = tf.Print(x, [x], "hello")
_ = tf.identity(x_, name="output")
onnx_graph = tf2onnx.tfonnx.process_tf_graph(sess.graph,
custom_op_handlers={"Print": (print_handler, ["Identity", "mode"])},
extra_opset=[helper.make_opsetid(_TENSORFLOW_DOMAIN, 1)],
input_names=["input:0"],
output_names=["output:0"])
model_proto = onnx_graph.make_model("test")
with open("/tmp/model.onnx", "wb") as f:
f.write(model_proto.SerializeToString())
tf2onnx 的工作原理
转换器需要解决一些问题:
- 转换 protobuf 格式。由于格式类似,这一步很直接。
- 需要将 TensorFlow 类型映射到其 ONNX 对应物。
- 对于很多操作,TensorFlow 将形状参数作为输入传递,而 ONNX 希望将其作为属性。由于我们使用冻结图,转换器将获取输入作为常量,转换为属性并删除原始输入。
- 许多情况下,TensorFlow 将操作组合成多个简单的操作。转换器需要识别这些操作的子图,切割子图并替换为其 ONNX 对应物。这可能会变得非常复杂,所以我们使用图匹配库来处理它。一个很好的例子是 tensorflow 的 transpose 操作。
- TensorFlow 的默认数据格式是 NHWC,而 ONNX 需要 NCHW。转换器将插入 transpose 操作来处理这个问题。
- 有些操作例如 relu6 在 ONNX 中不受支持,但转换器可以由其他 ONNX 操作组成。
- ONNX 后端是新的,并且实现尚不完整。对于一些操作,转换器生成的操作需要处理现有后端中的问题。
第一步 - 从冻结图开始
tf2onnx 从冻结图开始。由于上述第 3 项的原因。
第二步 - tf 向 onnx 的 protobuf 的 1:1 转换
tf2onnx 首先从 TensorFlow protobuf 格式到 ONNX protobuf 格式进行简单转换,而不关注单个操作。 我们这样做是为了可以使用 ONNX 图作为内部表示并围绕它编写辅助函数。 执行转换的代码在 tensorflow_to_onnx() 中。tensorflow_to_onnx() 将返回 ONNX 图和来自 TensorFlow 的形状信息字典。当处理单独的操作时,形状信息在某些情况下非常有用。 ONNX 图被包装在一个 Graph 对象中,图中的节点被包装在 Node 对象中,以允许在图上进行更容易的图操作。所有处理节点和图的代码都在 graph.py 中。
第三步 - 重写子图
在下一步中,我们在图上应用图匹配代码以重新编写 transpose 和 lstm 等操作的子图。例如请看 rewrite_transpose()。
第四步 - 处理单个操作
在第四步中,我们处理需要注意的单个操作。字典 _OPS_MAPPING 将 tensorflow 操作类型映射到用于处理操作的方法。最简单的情况是 direct_op(),操作可以照原样使用。每当可能时,我们尝试将操作分组到常见处理方法中,例如所有需要处理广播的操作都映射到 broadcast_op()。对于将 tensorflow 操作由多个 onnx 操作组成的操作,请参见 relu6_op()。
第五步 - 优化功能性 ONNX 图
然后我们尝试优化功能性 ONNX 图。例如我们删除不需要的操作, 尽可能地删除 transpose,去重常量,尽可能融合操作,...
第六步 - 最后处理
一旦所有操作都转换并优化,我们需要进行拓扑排序,因为 ONNX 需要它。process_tf_graph() 是一个负责上述所有步骤的方法。
扩展 tf2onnx
如果您想贡献并添加新的转换到 tf2onnx,流程如下:
1. 查看该操作是否适合现有的映射。如果适合,只需将其添加到 _OPS_MAPPING 中。
2. 如果新的操作需要额外处理,则启动一个新的映射函数。
3. 如果 tensorflow 操作由多个操作组成,考虑使用图重写。虽然这可能会稍微难一些,但它对复杂模式效果更好。
4. 在 tests/test_backend.py 中添加单元测试。单元测试主要是创建 tensorflow 图,运行并捕获输出,然后转换为 onnx,运行在 onnx 后端并比较 tensorflow 和 onnx 的结果。
5. 如果有使用新操作的预训练模型,考虑将它们添加到 test/run_pretrained_models.py 中。
## 许可证
[Apache 许可证 v2.0](LICENSE)