我正在关注诗人的张量流(tflite)教程: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为poets运行label_image脚本,这会产生这个错误:
KeyError:“名称'导入/输入'是指操作不在图中。”
寻找解决方案,我尝试将input_layer =“input”更改为input_layer =“Mul”,但这只会产生错误:
KeyError:“名称'import / Mul'是指不在图表中的操作。”
如果对我的错误有任何建议,我们将不胜感激。
答案 0 :(得分:0)
您是否尝试使用summarize_graph检查模型的潜在输入/输出节点名称?
答案 1 :(得分:0)
根据您的代码,您的input_layer名称为"前缀/输入"不是"输入"。更改为input_layer="prefix/input"
可以解决您的问题。