python

关注公众号 jb51net

关闭
首页 > 脚本专栏 > python > tensorflow转onnx

tensorflow转onnx的实现方法

作者:小广向前进

本文主要介绍了tensorflow转onnx的实现方法,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友们下面随着小编来一起学习学习吧

安装tf2onnx以及onnxruntime

pip install onnxruntime
pip install tf2onnx

tf 转为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

参数说明:

通过程序直接转:

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

到此这篇关于tensorflow转onnx的实现方法的文章就介绍到这了,更多相关tensorflow转onnx内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家!

您可能感兴趣的文章:
阅读全文