TensorFlow模型转换为ONNX格式的详细步骤与命令解析
发布时间：2026-10-12 |
3人浏览过
本文详细介绍将TensorFlow模型转换为ONNX格式的具体流程。首先通过Python代码冻结动态图生成PB文件，随后利用tf2onnx工具进行转换。重点解析命令行参数，特别是针对TF默认NHWC与ONNX NCHW布局差异的--inputs-as-nchw选项，并提供直接通过API转换Keras模型的方法，适用于模型部署场景。
执行安装命令以获取tf2onnx和onnxruntime依赖库。
pip install onnxruntime pip install tf2onnx
TensorFlow模型转换为ONNX格式主要包含以下步骤：
- 第一步是将TensorFlow动态图冻结，从而生成冻结后的PB模型文件。
- 第二步是使用tf2onnx工具将生成的PB文件转换为ONNX格式文件。
实现TensorFlow动态图冻结的具体代码如下：
def export_frozen_graph(model, model_dir, name_pb) :
f = tf.function(lambda x: model(inputs=x))
f = f.get_concrete_function(x=(tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype)))
frozen_func = convert_variables_to_constants_v2(f)
frozen_func.graph.as_graph_def()
print("-" * 50)
print("Frozen model inputs: ")
print(frozen_func.inputs)
print("Frozen model outputs: ")
print(frozen_func.outputs)
tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
logdir=model_dir,
name=name_pb,
as_text=False)在使用tf2onnx将PB文件转换为ONNX文件时，需在终端执行相关命令。需要注意的是，大多数TensorFlow模型的输入布局（layout）为NHWC，而ONNX模型的标准输入布局为NCHW。因此，建议在转换命令中添加--inputs-as-nchw选项以适配布局差异。其他可用选项可参考官方文档，具体运行命令示例如下：
python -m tf2onnx.convert --input yolo.pb --output model.onnx --outputs Identity:0,Identity_1:0,Identity_2:0 --inputs x:0 --inputs-as-nchw x:0 --opset 10
各参数含义说明如下：
- input：指定输入的PB模型文件路径。
- output：指定输出的ONNX文件名。
- inputs：指定输入层的名称，若存在多个输入层，各层名之间使用逗号分隔。
- outputs：指定输出层的名称，若存在多个输出层，各层名之间使用逗号分隔。
- --inputs-as-nchw：强制将输入层数据视为NCHW格式，使用时需明确指定对应的输入层名称。
- --opset：指定ONNX的操作集版本号。
此外，也可以通过程序API直接进行转换，示例代码如下：
tf2onnx.convert.from_keras(model, inputs_as_nchw=[model.inputs[0].name], output_path=model_filepath + 'yolo.onnx') --opset 10
至此，关于TensorFlow转换为ONNX格式的实现方法已介绍完毕。如需了解更多相关内容，可搜索本站过往文章或继续浏览后续文章，感谢支持。
本站声明
本文内容仅供参考，不代表本站观点。转载或引用请注明出处。如发现内容问题，请通过联系我们页面反馈。
