TensorFlow Lite 无法识别操作 VarHandleOp

Question

我正在尝试将 TF 模型转换为 TFLite。模型以 .pb格式保存，我使用以下代码对其进行了转换：

import os

import tensorflow as tf

from tensorflow.core.protobuf import meta_graph_pb2

export_dir = os.path.join('export_dir', '0')

if not os.path.exists('export_dir'):

    os.mkdir('export_dir')

tf.compat.v1.enable_control_flow_v2()

tf.compat.v1.enable_v2_tensorshape()

# I took this function from a tutorial on the TF website

def wrap_frozen_graph(graph_def, inputs, outputs):

    def _imports_graph_def():

        tf.compat.v1.import_graph_def(graph_def, name="")

    wrapped_import = tf.compat.v1.wrap_function(_imports_graph_def, [])

    import_graph = wrapped_import.graph

    return wrapped_import.prune(

            inputs, outputs)

graph_def = tf.compat.v1.GraphDef()

loaded = graph_def.ParseFromString(open(os.path.join(export_dir, 'saved_model.pb'),'rb').read())

concrete_func = wrap_frozen_graph(

        graph_def, inputs=['extern_data/placeholders/data/data:0', 'extern_data/placeholders/data/data_dim0_size:0'],

    outputs=['output/output_batch_major:0'])

concrete_func.inputs[0].set_shape([None, 50])

concrete_func.inputs[1].set_shape([None])

concrete_func.outputs[0].set_shape([None, 100])

converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])

converter.experimental_new_converter = True

converter.post_training_quantize=True

converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS,

                                               tf.lite.OpsSet.SELECT_TF_OPS]

converter.allow_custom_ops=True

tflite_model = converter.convert()

# Save the model.

if not os.path.exists('tflite'):

    os.mkdir('tflite')

output_model = os.path.join('tflite', 'model.tflite')

with open(output_model, 'wb') as f:

     f.write(tflite_model)

现在，我在代码中找不到任何内容VarHandleOp，我发现它实际上是在tensorflow中（https://www.tensorflow.org/api_docs/python/tf/raw_ops/VarHandleOp）。那么，为什么TFLite无法识别呢？

Answer 1

在模型转换的情况下提供最小的可重现示例当然很困难，但是这些问题将受益于更好的指针。如果你有一个“冻结图”（其中tf.Graph 变量已变成常量）（粗体突出部分是我的）。显然，您的图表包含VarHandleOp（这同样适用于Variable和VariableV2节点），并且根据此定义并未“冻结”。您的一般方法是有意义的，但您需要一个图表，其中包含节点形式的变量的实际训练值Const。您在训练时需要变量，但在推理时需要变量，并且应该将其烘焙到图中。TFLite 作为推理时间框架，不支持变量。你的其余想法似乎不错。TFLiteConverter.from_concrete_functions目前正好需要 1 concrete_function，但这就是您通过包装图表得到的结果。如果运气足够好，它可能会起作用。有一个实用程序tensorflow/python/tools/freeze_graph.py会尽力用从最新检查点文件中获取的常量替换 Graph.pb 中的变量。如果您查看其代码，可以使用保存的元图 ( checkpoint_name .meta ) 文件或将工具指向训练目录，从而消除大量猜测；另外，我认为提供模型目录是获得单个冻结图和分片模型的唯一方法。我注意到您在示例中input使用了 just 。tf.nest.map_structure(import_graph.as_graph_element, inputs)您可能有其他原因，但如果您这样做是因为as_graph_element抱怨数据类型/形状，则可以通过正确冻结图形来解决这个问题。从冻结图中获得的具体函数将对其输入形状和数据类型有一个很好的了解。一般来说，需要手动设置它们是出乎意料的，而且你这样做的事实对我来说似乎很奇怪（但我并不声称对 TF 的这个黑暗角落有广泛的经验）。map_structure有一个关键字参数来跳过检查。