简体   繁体   中英

Tensorflow TOCO python API

I'm following the tensorflow for poets (tflite) tutorial here: https://codelabs.developers.google.com/codelabs/tensorflow-for-poets-2-tflite/#3

I'm attempting to convert a customised graph from .pb to tflite using the python TOCO API : https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/toco/g3doc/python_api.md

The following code loads the retrained_graph.pb file, finds the input and output tensor, then calls toco_convert and writes the .tflite file.

    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)

This produces a test.tflite file. To check to see if it works, I run the label_image script from tf for poets, which produces this error:

KeyError: "The name 'import/input' refers to an Operation not in the graph."

Looking for solutions, I tried changing input_layer = "input" to input_layer = "Mul", but this only produces the error:

KeyError: "The name 'import/Mul' refers to an Operation not in the graph."

If there are any suggestions to what I am doing wrong they would be greatly appreciated.

您是否尝试过使用summary_graph检查模型的潜在输入/输出节点名称?

According to your code, your input_layer name is "prefix/input" not "input". Changing to input_layer="prefix/input" should solve your issues.

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM