tensorflow轉(zhuǎn)onnx如何實現(xiàn)

本文小編為大家詳細(xì)介紹“tensorflow轉(zhuǎn)onnx如何實現(xiàn)”，內(nèi)容詳細(xì)，步驟清晰，細(xì)節(jié)處理妥當(dāng)，希望這篇“tensorflow轉(zhuǎn)onnx如何實現(xiàn)”文章能幫助大家解決疑惑，下面跟著小編的思路慢慢深入，一起來學(xué)習(xí)新知識吧。

安裝tf2onnx以及onnxruntime

pip install onnxruntime
pip install tf2onnx

tf 轉(zhuǎn)為onnx步驟為如下：

• 將tf動態(tài)圖凍結(jié)，生成凍結(jié)后的pb文件

• 使用tf2onnx 將pb文件轉(zhuǎn)為onnx文件

將tf動態(tài)圖凍結(jié)使用如下代碼：

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文件轉(zhuǎn)為onnx文件時需要在終端執(zhí)行，需要指出的是大部分tf模型的輸入layout都是NHWC，而ONNX模型的輸入layout為NCHW，因此建議在轉(zhuǎn)換的時候加上--inputs-as-nchw這個選項，其他選項可以參考文檔，非常詳細(xì)，具體運(yùn)行命令如下：

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

參數(shù)說明：

• input 輸入的pb模型

• output 輸出的onnx文件名

• inputs 輸入層名字，有多個輸入時，中間用逗號隔開

• outputs 輸出層名字，有多個輸出時，中間用逗號隔開

• –inputs-as-nchw 將輸入作為nchw格式，注意加上輸入層名字

• –opset onnx 版本號

通過程序直接轉(zhuǎn)：

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

讀到這里，這篇“tensorflow轉(zhuǎn)onnx如何實現(xiàn)”文章已經(jīng)介紹完畢，想要掌握這篇文章的知識點還需要大家自己動手實踐使用過才能領(lǐng)會，如果想了解更多相關(guān)內(nèi)容的文章，歡迎關(guān)注億速云行業(yè)資訊頻道。