簡體   English   中英

Tensorflow TOCO python API

[英]Tensorflow TOCO python API

我在這里關注詩人(tflite)的tensorflow教程: https ://codelabs.developers.google.com/codelabs/tensorflow-for-poets-2-tflite/#3

我正在嘗試使用python TOCO API將自定義圖從.pb轉換為tflite: https : //github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/toco/g3doc/python_api.md

以下代碼加載retrained_graph.pb文件,找到輸入和輸出張量,然后調用toco_convert並寫入.tflite文件。

    import tensorflow as tf
     def load_graph(graph_filename):
     with tf.gfile.GFile(graph_filename, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())

     with tf.Graph().as_default() as graph:
        tf.import_graph_def(
          graph_def,
          input_map=None,
          return_elements=None,
          name="prefix",
          op_dict=None,
          producer_op_list=None
         )


    graph = load_graph("retrained_graph.pb")
    x = graph.get_tensor_by_name('prefix/input:0') #input tensor
    y = graph.get_tensor_by_name('prefix/final_result:0') #output tensor


    with tf.Session(graph=graph) as sess:
       tflite_model = tf.contrib.lite.toco_convert(sess.graph_def, [x], [y])
       open("test.tflite", "wb").write(tflite_model)

這將生成一個test.tflite文件。 要檢查它是否有效,我從tf運行詩人的label_image腳本,這會產生此錯誤:

KeyError:“名稱'import / input'指的是圖中未包含的操作。”

尋找解決方案時,我嘗試將input_layer =“ input”更改為input_layer =“ Mul”,但這只會產生錯誤:

KeyError:“名稱'import / Mul'指的是圖中未包含的操作”。

如果對我做錯的事情有任何建議,將不勝感激。

您是否嘗試過使用summary_graph檢查模型的潛在輸入/輸出節點名稱?

根據您的代碼,您的input_layer名稱是“前綴/輸入”而不是“輸入”。 更改為input_layer="prefix/input"應該可以解決您的問題。

暫無
暫無

聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.

 
粵ICP備18138465號  © 2020-2024 STACKOOM.COM