TensorFlow Lite 无法识别操作 VarHandleOp答案

【问题标题】：TensorFlow Lite does not recognize op VarHandleOpTensorFlow Lite 无法识别操作 VarHandleOp
【发布时间】：2020-10-07 08:00:46
【问题描述】：

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

但是，当我尝试将此模型与解释器一起使用时，我收到以下错误：

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) 中。那么，为什么 TFLite 无法识别呢？

【问题讨论】：

.pb 中的 TFv1 是什么意思？ TFv2 == TF 2？您使用什么 TF 版本来运行发布的脚本？为什么不总是使用 TF 2？
TFLite 不需要冻结图吗？ IE。将所有变量转换为常量。 VarHandleOp 用于变量。 TFLite 是否支持变量？

标签： python tensorflow tensorflow2.0 tensorflow-lite

【解决方案1】：

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

您所指的教程可能来自section on migrating from TF1 to TF2，特别是处理原始图形文件的部分。至关重要的一点是

如果你有一个“冻结图”（tf.Graph变量已转换为常量）

（粗体部分是我的）。显然，您的图表包含VarHandleOp（同样适用于Variable 和VariableV2 节点），并且根据这个定义没有“冻结”。您的一般方法是有意义的，但您需要一个图表，其中包含 Const 节点形式的变量的实际训练值。您在训练时需要变量，但在推理时需要变量，并且应该将其烘焙到图中。 TFLite 作为推理时框架，不支持变量。

你的想法的其余部分似乎很好。 TFLiteConverter.from_concrete_functions 当前只需要一个 concrete_function，但这是你从包装图表中得到的。如果运气够好，它可能会奏效。

有一个实用程序tensorflow/python/tools/freeze_graph.py 会尽力用最新检查点文件中的常量替换 Graph.pb 中的变量。如果您查看它的代码，无论是使用保存的元图 (checkpoint_name.meta) 文件还是将工具指向训练目录，都可以消除很多猜测；另外，我认为提供模型目录是使单个冻结图成为分片模型的唯一方法。

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

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

【讨论】：