TFLite 中缺少广播的解决方法答案

【问题标题】：Workaround for lack of broadcast in TFLiteTFLite 中缺少广播的解决方法
【发布时间】：2019-09-23 20:11:46
【问题描述】：

我想运行一个需要我生成 3d 输出的 TFLite 模型（示例代码是生成错误的最小示例）。有没有等价于gather_nd的张量流不降维？

我已经尝试在文档中查找我能想到的相关功能，但没有找到好的选择。

import tensorflow.compat.v1 as tf
import numpy as np

tf.disable_v2_behavior()
initial_input = tf.placeholder(dtype=tf.float32, shape=(None,5,1024))
cap_i = tf.gather_nd(initial_input, [[0,1]]) #[0,2],[0,3],[0,4],[0,5]
cap_i_broadcast = tf.broadcast_to(cap_i, [1,5,1024])
cap_iT = tf.transpose(cap_i_broadcast, perm=[0,2,1])

sess = tf.Session()
sess.run(tf.global_variables_initializer())
tf.io.write_graph(sess.graph_def, '', 'train.pbtxt')
converter = tf.lite.TFLiteConverter.from_session(sess, [initial_input], [cap_iT])
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]
tflite_model = converter.convert()
open('converted_model.tflite', "wb").write(tflite_model)
sess.close()

模型中的某些运算符不受标准 TensorFlow Lite 运行时支持，也无法被 TensorFlow 识别。如果您有自定义实现，您可以使用 --allow_custom_ops 禁用此错误，或者在调用 tf.lite.TFLiteConverter() 时设置 allow_custom_ops=True。以下是您正在使用的内置运算符列表：GATHER_ND、TRANSPOSE。以下是您需要自定义实现的运算符列表：BroadcastTo。

【问题讨论】：

标签： tensorflow

【解决方案1】：

以下代码有一个解决方案，使用带降维的跨步切片，然后重新整形以恢复正确的维度。

import tensorflow.compat.v1 as tf
import numpy as np

tf.disable_v2_behavior()
initial_input = tf.placeholder(dtype=tf.float32, shape=(None,5,1024))
cap_i = tf.strided_slice(initial_input, [0,0,0], [0,5,1024], [1,1,1], 
shrink_axis_mask=1)
cap_i_reshaped =tf.reshape(cap_i,[1,5,1024])
cap_iT = tf.transpose(cap_i_reshaped, perm=[0,2,1])

sess = tf.Session()
sess.run(tf.global_variables_initializer())
tf.io.write_graph(sess.graph_def, '', 'train.pbtxt')
converter = tf.lite.TFLiteConverter.from_session(sess, [initial_input], 
[cap_iT])
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, 
tf.lite.OpsSet.SELECT_TF_OPS]
tflite_model = converter.convert()
open('converted_model.tflite', "wb").write(tflite_model)
sess.close()

以前认为 TFLite 支持 slice，但现在只有 strided_slice 支持。

【讨论】：