tensorflow转onnx如何实现

本文小编为大家详细介绍“tensorflow转onnx如何实现”，内容详细，步骤清晰，细节处理妥当，希望这篇“tensorflow转onnx如何实现”文章能帮助大家解决疑惑，下面跟着小编的思路慢慢深入，一起来学习新知识吧。

安装tf2onnx以及onnxruntime

pip install onnxruntime
pip install tf2onnx

tf 转为onnx步骤为如下：

• 将tf动态图冻结，生成冻结后的pb文件

• 使用tf2onnx 将pb文件转为onnx文件

将tf动态图冻结使用如下代码：

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文件时需要在终端执行，需要指出的是大部分tf模型的输入layout都是NHWC，而ONNX模型的输入layout为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 版本号

通过程序直接转：

tf2onnx.convert.from_keras(model, inputs_as_nchw=[model.inputs[0].name], output_path=model_filepath + 'yolo.onnx') --opset 10

读到这里，这篇“tensorflow转onnx如何实现”文章已经介绍完毕，想要掌握这篇文章的知识点还需要大家自己动手实践使用过才能领会，如果想了解更多相关内容的文章，欢迎关注亿速云行业资讯频道。