概述
TensorFlow Lite 支持将 TensorFlow RNN 模型转换为 TensorFlow Lite 的融合 LSTM 运算。融合运算旨在最大限度地提高其底层内核实现的性能，并提供更高级别的接口来定义复杂的转换，例如量化。
由于 TensorFlow 中存在许多 RNN API 变体，因此我们的方法是双重的
- 为 Keras LSTM 等标准 TensorFlow RNN API 提供原生支持。这是推荐的选项。
- 为用户定义的 RNN 实现提供一个接口进入转换基础设施，以插入并转换为 TensorFlow Lite。我们提供了一些使用 lingvo 的此类转换的开箱即用示例，例如 LSTMCellSimple 和 LayerNormalizedLSTMCellSimple RNN 接口。
转换器 API
该功能是 TensorFlow 2.3 版本的一部分。它也可以通过 tf-nightly pip 或从头开始获得。
当通过 SavedModel 或直接从 Keras 模型转换为 TensorFlow Lite 时，此转换功能可用。请参阅示例用法。
从保存的模型
# build a saved model. Here concrete_function is the exported function
# corresponding to the TensorFlow model containing one or more
# Keras LSTM layers.
saved_model, saved_model_dir = build_saved_model_lstm(...)
saved_model.save(saved_model_dir, save_format="tf", signatures=concrete_func)
# Convert the model.
converter = TFLiteConverter.from_saved_model(saved_model_dir)
tflite_model = converter.convert()
从 Keras 模型
# build a Keras model
keras_model = build_keras_lstm(...)
# Convert the model.
converter = TFLiteConverter.from_keras_model(keras_model)
tflite_model = converter.convert()
示例
Keras LSTM 到 TensorFlow Lite Colab 说明了 TensorFlow Lite 解释器的端到端用法。
支持的 TensorFlow RNN API
Keras LSTM 转换（推荐）
我们支持将 Keras LSTM 开箱即用地转换为 TensorFlow Lite。有关其工作原理的详细信息，请参阅 Keras LSTM 接口 以及转换逻辑 此处。
同样重要的是要强调 TensorFlow Lite 的 LSTM 与 Keras 运算定义的契约
- 输入张量的维度 0 是批次大小。
- 循环权重张量的维度 0 是输出数量。
- 权重和循环内核张量被转置。
- 转置后的权重、转置后的循环内核和**偏差**张量沿维度 0 分割成 4 个大小相等的张量。它们分别对应于**输入门、遗忘门、单元和输出门**。
Keras LSTM 变体
时间主导
用户可以选择时间主导或非时间主导。Keras LSTM 在函数定义属性中添加了一个时间主导属性。对于单向序列 LSTM，我们可以简单地映射到 unidirecional_sequence_lstm 的 时间主导属性。
双向 LSTM
双向 LSTM 可以使用两个 Keras LSTM 层实现，一个用于前向，一个用于后向，请参阅示例 此处。一旦我们看到 go_backward 属性，我们就会将其识别为后向 LSTM，然后我们将前向和后向 LSTM 组合在一起。**这是未来的工作。**目前，这会在 TensorFlow Lite 模型中创建两个 UnidirectionalSequenceLSTM 操作。
用户定义的 LSTM 转换示例
TensorFlow Lite 还提供了一种方法来转换用户定义的 LSTM 实现。这里我们以 Lingvo 的 LSTM 为例来说明如何实现。有关详细信息，请参阅 lingvo.LSTMCellSimple 接口 和转换逻辑 此处。我们还提供另一个 Lingvo 的 LSTM 定义的示例 lingvo.LayerNormalizedLSTMCellSimple 接口 及其转换逻辑 此处。
将“自带 TensorFlow RNN”引入 TensorFlow Lite
如果用户的 RNN 接口与标准支持的接口不同，则有以下两种选择
**选项 1：**在 TensorFlow python 中编写适配器代码，将 RNN 接口适配到 Keras RNN 接口。这意味着一个带有 tf_implements 注释 的 tf.function，该注释位于生成的 RNN 接口的函数上，该函数与 Keras LSTM 层生成的函数相同。之后，用于 Keras LSTM 的相同转换 API 将起作用。
**选项 2：**如果上述方法不可行（例如，Keras LSTM 缺少 TensorFlow Lite 的融合 LSTM 操作当前公开的一些功能，例如层归一化），则通过编写自定义转换代码扩展 TensorFlow Lite 转换器，并将其插入到 prepare-composite-functions MLIR-pass 中 此处。函数的接口应被视为 API 合同，应包含转换为融合 TensorFlow Lite LSTM 操作所需的参数 - 即输入、偏差、权重、投影、层归一化等。最好将作为参数传递给该函数的张量具有已知的秩（即 MLIR 中的 RankedTensorType）。这使得编写可以将这些张量视为 RankedTensorType 的转换代码变得容易得多，并有助于将它们转换为与融合 TensorFlow Lite 操作的操作数相对应的秩张量。
Lingvo 的 LSTMCellSimple 到 TensorFlow Lite 转换的完整示例。
Lingvo 中的 LSTMCellSimple 定义 此处。使用此 LSTM 单元训练的模型可以转换为 TensorFlow Lite，如下所示
- 将所有使用 LSTMCellSimple 的地方包装在一个带有 tf_implements 注释的 tf.function 中，该注释标记为这样（例如，lingvo.LSTMCellSimple 将是一个很好的注释名称）。确保生成的 tf.function 与转换代码中预期的函数接口匹配。这是在添加注释的模型作者和转换代码之间的契约。
扩展 prepare-composite-functions 传递，将自定义复合操作插入到 TensorFlow Lite 融合 LSTM 操作转换中。请参阅 LSTMCellSimple 转换代码。
转换契约
**权重**和**投影**张量被转置。
通过切片转置后的权重张量来提取**{输入、循环}** 到 **{单元、输入门、遗忘门、输出门}**。
通过切片偏差张量来提取**{偏差}** 到 **{单元、输入门、遗忘门、输出门}**。
通过切片转置后的投影张量来提取**投影**。
类似的转换已为 LayerNormalizedLSTMCellSimple 编写。
TensorFlow Lite 转换基础设施的其余部分，包括所有定义的 MLIR 传递 以及最终导出到 TensorFlow Lite flatbuffer，都可以重复使用。
已知问题/限制
- 目前仅支持转换无状态 Keras LSTM（Keras 中的默认行为）。有状态 Keras LSTM 转换是未来的工作。
- 仍然可以使用底层的无状态 Keras LSTM 层来模拟有状态 Keras LSTM 层，并在用户程序中显式管理状态。这样的 TensorFlow 程序仍然可以使用此处描述的功能转换为 TensorFlow Lite。
- 双向 LSTM 目前在 TensorFlow Lite 中被建模为两个 UnidirectionalSequenceLSTM 操作。这将被替换为单个 BidirectionalSequenceLSTM 操作。