如何让TensorFlow Lite转换模型后输出标准LSTM层而非单向序列LSTM

问题原因 #

TFLite的新转换器默认会将Keras原生LSTM层自动融合为优化后的单向序列LSTM（UNIDIRECTIONNAL_SEQUENCE_LSTM）算子，用于提升端侧推理性能。如果需要保留标准LSTM算子结构，可通过以下两种方案解决：

方案1：禁用LSTM算子融合（优先使用） #

该方案生成的模型仅依赖TFLite内置算子，无需额外依赖，体积小推理速度快。修改转换代码如下：

converter = tf.lite.TFLiteConverter.from_saved_model("mnist_lstm_model")
converter.experimental_new_converter = True
# 新增配置，禁用LSTM算子融合逻辑
converter.experimental_disable_lstm_fusion = True
tflite_model = converter.convert()

# 保存转换后的模型
with open("mnist_lstm_std.tflite", "wb") as f:
    f.write(tflite_model)

注意：该参数在TensorFlow 2.4及以上版本生效，如果你的TF版本低于2.4建议先升级到稳定新版本。

方案2：使用TF Select算子保留原生LSTM结构 #

如果方案1无法满足需求（比如需要完全对齐训练时的LSTM逻辑），可以通过支持TensorFlow原生算子的方式转换，修改代码如下：

converter = tf.lite.TFLiteConverter.from_saved_model("mnist_lstm_model")
converter.experimental_new_converter = True
# 配置同时支持TFLite内置算子和TensorFlow原生算子
converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS,
    tf.lite.OpsSet.SELECT_TF_OPS
]
# 关闭张量列表降级逻辑，避免LSTM被拆分融合
converter._experimental_lower_tensor_list_ops = False
tflite_model = converter.convert()

# 保存转换后的模型
with open("mnist_lstm_original.tflite", "wb") as f:
    f.write(tflite_model)

注意：该方案生成的模型需要部署环境支持TF Select算子库，整体体积会比方案1的模型更大。

内容的提问来源于stack exchange，提问作者Lucas Kassianow