在 tensorflow 转换处理期间如何分配或修复特殊操作？

Question

我在将 tensorflow 会话转换为 tensorflow lite 模型时遇到问题。
问题是 TFLiteConverter 强制使用 FlexFloorDiv 替换原始函数（tf.floordiv），并且 tflite 不支持 FlexFloorDiv。
但是我已经检查了 tflite 的 API 列表，我发现 tflite 支持 floor_div (https://tensorflow.google.cn/mlir/tfl_ops#tflfloor_div_mlirtflfloordivop).
我想要一个tflite原本支持的转换，换句话说，我只想使用'tf.lite.OpsSet.TFLITE_BUILTINS'.

import tensorflow as tf
if tf.__version__ > '2':
    import tensorflow.compat.v1 as tf  
    tf.disable_v2_behavior()


tf.compat.v1.reset_default_graph()
inputs = tf.compat.v1.placeholder(tf.float32, shape=[20, 4])

maxloc = tf.argmax(input=inputs, axis=0)
inputs_shape = tf.shape(input=inputs)

# Here is the error code
# the converter will use FlexFloorDiv not floor_div
out = tf.floordiv(maxloc, tf.cast(inputs_shape[1], tf.int64))

sess = tf.compat.v1.Session()
sess.run(tf.compat.v1.global_variables_initializer())
sess.run(tf.compat.v1.local_variables_initializer())
converter = tf.lite.TFLiteConverter.from_session(sess, [inputs], [out])
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]
# converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS] # can not convert.
converter.allow_custom_ops = False
converter.experimental_new_converter = True
tflite_model = converter.convert()

输出为：

2022-01-21 18:01:40.286574: W tensorflow/compiler/mlir/lite/flatbuffer_export.cc:1891] TFLite interpreter needs to link Flex delegate in order to run the model since it contains the following Select TFop(s):
Flex ops: FlexFloorDiv
Details:
        tf.FloorDiv(tensor<4xi64>, tensor<i64>) -> (tensor<4xi64>) : {device = ""}
See instructions: https://www.tensorflow.org/lite/guide/ops_select
2022-01-21 18:01:40.286641: I tensorflow/compiler/mlir/lite/flatbuffer_export.cc:1962] Estimated count of arithmetic ops: 0  ops, equivalently 0  MACs

Answer 1

可能tensorflow-lite对int32\int64类型的div运算符不友好，会转为tf.RealDiv或tf.FlexDiv，不支持原始 tflite.
为了获得我的目标，这里我使用 float32 div 而不是 int64 div，代码如下：

maxloc = tf.argmax(input=inputs, axis=0)
inputs_shape = tf.shape(input=inputs)
out = tf.cast(out, tf.float32) / tf.cast(inputs_shape [1], tf.float32)
# use cutoff feature when from float to int, we can get floor operator.  
out = tf.cast(out , tf.int64)

在 tensorflow 转换处理期间如何分配或修复特殊操作？

How can I assign or fix special op during convert processing in tensorflow?

tensorflow

tensorflow-lite