概览
TensorFlow Lite 支持将 TensorFlow RNN 模型转换为 TensorFlow Lite 的融合 LSTM 操作。融合操作旨在最大限度地提高其底层内核实现的性能，并提供更高级别的接口来定义量化等复杂转换。
由于 TensorFlow 中 RNN API 有许多变体，因此我们采取了两种方法：
- 提供对标准 TensorFlow RNN API 的原生支持，例如 Keras LSTM。 此为推荐选项。
- 为用户定义的 RNN 实现提供用于转换基础架构的接口，以便插入并转换为 TensorFlow Lite。我们使用 lingvo 的 LSTMCellSimple 和 LayerNormalizedLSTMCellSimple RNN 接口提供此类转换的几个现成示例。
转换器 API
此功能是 TensorFlow 2.3 版本的一部分。也可以通过 tf-nightly pip 或 head 运行。
通过 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 RNNs API
Keras LSTM 转换（推荐）
我们支持开箱即用将 Keras LSTM 转换为 TensorFlow Lite。如需详细了解其工作原理，请参阅 Keras LSTM 接口 以及此处的转换逻辑。
另外，重点强调关于 Keras 操作定义的 TensorFlow Lite 的 LSTM 协定：
- input 张量的维度 0 是批次大小。
- recurrent_weight 张量的维度 0 是输出的数量。
- weight 和 recurrent_kernel 张量进行了转置。
- 转置权重、转置 recurrent_kernel 和 bias 张量在维度 0 上被拆分为 4 个大小相等的张量。这三个对应关系对应于输入门、遗忘门、单元和输出门。
Keras LSTM 变体
时间专业
用户可以选择“主要”或不“主要”。Keras LSTM 在函数 def 属性中添加了一个时间主要属性。对于单向序列 LSTM，我们只需映射到 unidirecional_sequence_lstm 的时间主属性。
双向 LSTM
可以使用两个 Keras LSTM 层来实现双向 LSTM，一个用于前向层，一个用于后向。如需查看示例，请点击此处。 看到 go_backward 属性后，就会将其识别为后向 LSTM，接着将前向和后向 LSTM 组合在一起。这是未来的工作。目前，这会在 TensorFlow Lite 模型中创建两个 UniDirectionSequenceLSTM 操作。
用户定义的 LSTM 转化示例
TensorFlow Lite 还提供了一种转换用户定义的 LSTM 实现的方法。在这里，我们以 Lingvo 的 LSTM 为例来说明如何实现该机制。如需了解详情，请参阅 lingvo.LSTMCellSimple 接口和此处的转换逻辑。我们还在 lingvo.LayerNormalizedLSTMCellSimple 接口中提供 Lingvo 的另一个 LSTM 定义示例及其转换逻辑，详见此处。
在 TensorFlow Lite 中“使用自己的 TensorFlow RNN”
如果用户的 RNN 接口与支持的标准 RNN 接口不同，则有以下几种选择：
方式 1：在 TensorFlow Python 中编写适配器代码，以使 RNN 接口适应 Keras RNN 接口。这意味着在生成的 RNN 接口的函数上有一个带有 tf_implements 注解的 tf.function，该函数与 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，如下所示：
- 使用带有相应标签的 tf_implements 注解将 LSTMCellSimple 的所有用例封装在 tf.function 中（例如，在此处使用 lingvo.LSTMCellSimple 作为一个不错的注解名称）。请确保生成的 tf.function 与转换代码中预期的函数的接口相匹配。这是添加注释的模型作者与转化代码之间的协定。
扩展 prepare-composite-functions 传递，以插入 TensorFlow Lite 融合 LSTM 运算的自定义复合运算。请参阅 LSTMCellSimple 转换代码。
转化协定：
Weight 和 projection 张量进行转置，
通过分割转置权重张量来提取到 {cell, input gate, forgot gate, output gate} 的 {input, recurrent}。
通过分割偏差张量提取 {bias} 到 {cell, input gate, forgot gate, output gate}。
projection 通过切片转置投影张量来提取。
类似的转换是针对 LayerNormalizedLSTMCellSimple 编写的。
TensorFlow Lite 转换基础架构的其余部分（包括已定义的所有 MLIR 卡券以及最终导出到 TensorFlow Lite 平板缓冲区）可以重复使用。
已知问题/限制
- 目前仅支持转换无状态 Keras LSTM（Keras 中的默认行为）。有状态 Keras LSTM 转换是未来的工作。
- 仍然可以使用底层无状态 Keras LSTM 层对有状态 Keras LSTM 层建模，并在用户程序中显式管理状态。此类 TensorFlow 程序仍然可以使用此处介绍的功能转换为 TensorFlow Lite。
- 双向 LSTM 目前在 TensorFlow Lite 中建模为两个单向序列 LSTM 指令。此函数将替换为单个 BidirectionalSequenceLSTM 操作。