Tensorflow 使用 python 将 pb 文件转换为 TFLITE

Tensorflow Convert pb file to TFLITE using python

我有一个模型在训练后保存为 pb 文件,我想使用 tensorflow mobile,使用 TFLITE 文件很重要。 问题是我在谷歌搜索转换器后发现的大多数示例都是终端或 cmd 上的命令。 您能否与我分享一个使用 python 代码转换为 tflite 文件的示例?

可以直接在python中直接转成tflite。你必须freeze the graph and use toco_convert。它需要在调用 API 之前确定输入和输出的名称和形状,就像在命令行情况下一样。

示例代码片段

documentation 复制,其中 "frozen"(无变量)图被定义为代码的一部分:

import tensorflow as tf

img = tf.placeholder(name="img", dtype=tf.float32, shape=(1, 64, 64, 3))
val = img + tf.constant([1., 2., 3.]) + tf.constant([1., 4., 4.])
out = tf.identity(val, name="out")
with tf.Session() as sess:
  tflite_model = tf.contrib.lite.toco_convert(sess.graph_def, [img], [out])
  open("test.tflite", "wb").write(tflite_model)

在上面的示例中,由于没有变量,因此没有冻结图步骤。如果你有变量和 运行 toco 没有冻结图,即首先将这些变量转换为常量,那么 toco 会报错!

如果你冻结了 graphdef 并且知道输入和输出

那你就不需要session了。可以直接调用toco API:

path_to_frozen_graphdef_pb = '...'
input_tensors = [...]
output_tensors = [...]
frozen_graph_def = tf.GraphDef()
with open(path_to_frozen_graphdef_pb, 'rb') as f:
  frozen_graph_def.ParseFromString(f.read())
tflite_model = tf.contrib.lite.toco_convert(frozen_graph_def, input_tensors, output_tensors)

如果你有非冻结的 graphdef 并且知道输入和输出

然后你必须加载会话并在调用 toco 之前先冻结图形:

path_to_graphdef_pb = '...'
g = tf.GraphDef()
with open(path_to_graphdef_pb, 'rb') as f:
  g.ParseFromString(f.read())
output_node_names = ["..."]
input_tensors = [..]
output_tensors = [...]

with tf.Session(graph=g) as sess:
  frozen_graph_def = tf.graph_util.convert_variables_to_constants(
      sess, sess.graph_def, output_node_names)
# Note here we are passing frozen_graph_def obtained in the previous step to toco.
tflite_model = tf.contrib.lite.toco_convert(frozen_graph_def, input_tensors, output_tensors)

如果您不知道图形的输入/输出

如果您没有定义图表,可能会发生这种情况,例如。您从某个地方下载了图表,或者使用了像 tf.estimators 这样的高级 API 来隐藏图表。在这种情况下,您需要在调用 toco 之前加载图形并四处查看以找出输入和输出。请参阅我对 .

的回答
converter = tf.contrib.lite.TFLiteConverter.from_frozen_graph(
    frozen_model_filename, INPUT_NODE, OUTPUT_NODE)

tflite_model = converter.convert()
open(TFLITE_OUTPUT_FILE, "wb").write(tflite_model)

INPUT_NODE和OUTPUT_NODE分别是输入和输出的名称列表。

在此TF example之后,您可以传递“--Saved_model_dir”参数以将saved_model.pb和变量文件夹导出到[=37之前的某个目录(none现有目录) =] retrain.py 脚本:

python retrain.py ...... --saved_model_dir /home/..../export

为了将您的模型转换为 tflite,您需要使用以下行:

convert_saved_model.convert(saved_model_dir='/home/.../export',output_arrays="final_result",output_tflite='/home/.../export/graph.tflite')

注意:需要导入convert_saved_model:

from tensorflow.contrib.lite.python import convert_saved_model

请记住,您可以通过两种方式转换为 tflite:

但是 最简单的方法是使用变量导出 saved_model.pb 以防你想避免使用像 Bazel 这样的构建工具。

这对我有用:(SSD_InceptionV2 模型)

  1. 完成训练后。我使用了 object_detection 文件夹中的 model_main.py。 TFv1.11
  2. 将图表导出为 TFLITE:
python /tensorflow/models/research/object_detection/export_tflite_ssd_graph.py

--pipeline_config_path annotations/ssd_inception_v2_coco.config 
--trained_checkpoint_prefix trained-inference-graphs/inference_graph_v7.pb/model.ckpt 
--output_directory trained-inference-graphs/inference_graph_v7.pb/tflite 
--max_detections 3
  1. 这会生成一个 .pb 文件,因此您可以像这样从中生成 tflite 文件:
tflite_convert 
--output_file=test.tflite 
--graph_def_file=tflite_graph.pb 
--input_arrays=normalized_input_image_tensor 
--output_arrays='TFLite_Detection_PostProcess','TFLite_Detection_PostProcess:1','TFLite_Detection_PostProcess:2','TFLite_Detection_PostProcess:3'

--input_shape=1,300,300,3 
--allow_custom_ops

现在 inputs/outputs 我不是 100 确定如何得到这个但是这段代码之前帮助我:

import tensorflow as tf
frozen='/tensorflow/mobilenets/mobilenet_v1_1.0_224.pb'
gf = tf.GraphDef()
gf.ParseFromString(open(frozen,'rb').read())
[n.name + '=>' +  n.op for n in gf.node if n.op in ( 'Softmax','Placeholder')]    
[n.name + '=>' +  n.op for n in gf.node if n.op in ( 'Softmax','Mul')]