TensorFlow Lite 无法识别操作 VarHandleOp

我正在尝试将 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)

但是，当我尝试将 intepretere 与该模型一起使用时，出现以下错误：

INFO: TfLiteFlexDelegate delegate: 8 nodes delegated out of 970 nodes with 3 partitions.

INFO: TfLiteFlexDelegate delegate: 0 nodes delegated out of 4 nodes with 0 partitions.

INFO: TfLiteFlexDelegate delegate: 3 nodes delegated out of 946 nodes with 1 partitions.

INFO: TfLiteFlexDelegate delegate: 0 nodes delegated out of 1 nodes with 0 partitions.

INFO: TfLiteFlexDelegate delegate: 3 nodes delegated out of 16 nodes with 2 partitions.

Traceback (most recent call last):
  File "/path/to/tflite_interpreter.py", line 9, in <module>
    interpreter.allocate_tensors()
  File "/path/to/lib/python3.6/site-packages/tensorflow/lite/python/interpreter.py", line 243, in allocate_tensors
    return self._interpreter.AllocateTensors()
RuntimeError: Encountered unresolved custom op: VarHandleOp.Node number 0 (VarHandleOp) failed to prepare.

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

正如 SO 指南所建议的那样，在模型转换的情况下提供最小的可重现示例当然很困难，但这些问题将受益于更好的指针。例如，与其说“我从 TF 网站上的教程中获取了这个函数”，不如提供教程的链接是一个更好的主意。 TF 网站非常庞大。

您所指的教程可能来自关于从 TF1 迁移到 TF2 的部分 https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt，特别是处理原始图形文件的部分。最重要的一点是

如果你有一个“冻结图”（atf.Graph 其中变量已转换为常量)

（粗体突出部分是我的）。显然，你的图表包含VarHandleOp（这同样适用于Variable and VariableV2节点），并且根据此定义不被“冻结”。您的一般方法是有意义的，但是您需要一个图表，其中包含以下形式的变量的实际训练值Const https://www.tensorflow.org/api_docs/python/tf/raw_ops/Const节点。您在训练时需要变量，但在推理时需要变量，并且应该将其烘焙到图中。 TFLite 作为推理时间框架，不支持变量。

你的其余想法似乎不错。TFLiteConverter.from_concrete_functions目前正好需要一个concrete_function，但这就是您通过包装图表得到的结果。如果运气足够好，它可能会起作用。

有一个实用程序tensorflow/python/tools/freeze_graph.py https://github.com/tensorflow/tensorflow/blob/v2.3.1/tensorflow/python/tools/freeze_graph.py尽力用从最新检查点文件中获取的常量替换 Graph.pb 中的变量。如果您查看其代码，可以使用保存的元图（检查点名称.meta）文件或将工具指向训练目录可以消除很多猜测；另外，我认为提供模型目录是获得单个冻结图和分片模型的唯一方法。

我注意到你只使用input代替tf.nest.map_structure(import_graph.as_graph_element, inputs)在示例中。你可能还有其他原因，但如果你这样做是因为as_graph_element抱怨数据类型/形状，这可能通过正确冻结图表来解决。从冻结图中获得的具体函数将对其输入形状和数据类型有一个很好的了解。一般来说，需要手动设置它们是出乎意料的，而且你这样做的事实对我来说似乎很奇怪（但我并不声称对 TF 的这个黑暗角落有广泛的经验）。

map_structure有一个关键字参数来跳过检查。