概览
LiteRT 支持将 TensorFlow RNN 模型转换为 LiteRT 模型 Fused LSTM 操作。融合操作旨在最大限度地提高 底层内核实现,并提供更高级别的 API, 接口来定义量化化等复杂的转换。
由于 TensorFlow 中有许多 RNN API 的变体,我们的方法 双层:
- 提供对标准 TensorFlow RNN API 的原生支持,例如 Keras LSTM。 此为推荐选项。
- 提供连接 转换基础架构的 用户定义的 RNN 实现,用于插入并转换为 LiteRT。我们在这里提供了几个现成的示例, 使用 lingvo 的 LSTMCellSimple 和 LayerNormalizedLSTMCellSimple RNN 接口。
转换器 API
该功能是 TensorFlow 2.3 版本的一部分。您还可以通过 tf-nightly pip 或从头创建。
转换为 LiteRT 后,此转换功能可用 从 Keras 模型导入,或直接从 Keras 模型导入,查看用法示例。
来自已保存的模型
# 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 到 LiteRT Colab 展示了 LiteRT 解释器的端到端用法。
支持 TensorFlow RNN API
Keras LSTM 转换(推荐)
我们支持开箱即用 Keras LSTM 到 LiteRT 的转换。对于 请参阅 Keras LSTM 接口 和转化逻辑 此处。
还有一点也很重要,那就是强调 LiteRT 的 LSTM 合同 添加到 Keras 运算定义:
- input 张量的维度 0 是批次大小。
- recurrent_weight 张量的维度 0 是 输出。
- weight 和 recurrent_kernel 张量经过转置。
- 转置权重、转置 recurrent_kernel 和 bias 张量 沿着维度 0 拆分为 4 个大小相等的张量。它们对应于 输入门、遗忘门、单元和输出门。
Keras LSTM 变体
时间专业
用户可以选择以时间为主,也可以选择不以时间为主。Keras LSTM 将 属性。对于单向序列 LSTM, 可以简单地映射到 unidirecional_sequence_lstm 中。 time primary 属性。
双向 LSTM
双向 LSTM 可以通过两个 Keras LSTM 层来实现,其中一个用于 分别表示前进和后退,请参阅示例 此处。 看到 go_backward 属性后,就会将其识别为后向 LSTM,然后 我们将前进和后退反向 LSTM 组合而成。这是未来的工作。目前, 这会在 LiteRT 模型。
用户定义的 LSTM 转化示例
LiteRT 还提供了一种方法来转换用户定义的 LSTM 实现。这里我们以 Lingvo 的 LSTM 为例 。如需了解详情,请参阅 lingvo.LSTMCellSimple 接口 以及转换逻辑 此处。 我们还举例说明了 Lingvo 的其他 LSTM 定义 lingvo.LayerNormalizedLSTMCellSimple 接口 及其转换逻辑 此处。
在 LiteRT 中“引入您自己的 TensorFlow RNN”
如果用户的 RNN 接口不同于支持的标准接口, 您可以采用以下两种方式:
方式 1:在 TensorFlow Python 中编写适配器代码以适应 RNN 接口 Keras RNN 接口,这意味着,具有如下特征的 tf.function tf_implements 注解 生成的 RNN 接口的函数,该函数与 Keras LSTM 层。在此之后,Keras LSTM 所用的转换 API 有效。
选项 2:如果不能做到以上所述(例如,Keras LSTM 缺少一些 目前由 LiteRT 的融合 LSTM 操作公开的功能,例如 层归一化),然后编写 并将自定义转化代码插入到 prepare-composite-functions MLIR-pass 此处。 函数的接口应被视为 API 协定, 包含转换为融合 LiteRT LSTM 所需的参数 运算 - 即输入、偏差、权重、投影、层归一化等。 最好让作为参数传递给该函数的张量具有已知 rank(即 MLIR 中的 RankedTensorType)。这样,在撰写代码时 转换代码,可以将这些张量假设为 RankedTensorType, 将其转换为与 fused LiteRT 对应的已排序张量 运算符的操作数。
此类转换流程的一个完整示例是 Lingvo 的 LSTMCellSimple LiteRT 转换。
Lingvo 中的 LSTMCellSimple 已定义 此处。 使用此 LSTM 单元训练的模型可以转换为 LiteRT,如下所示 如下:
- 使用 tf_implements 将 LSTMCellSimple 的所有使用过程封装在 tf.function 中 带有该标记的注释(例如 lingvo.LSTMCellSimple 应为 注释名称)。确保生成的 tf.function 与转化代码中预期函数的接口相符。这个 是添加了注解的模型作者与 转化代码。
扩展 prepare-composite-functions 传递以插入自定义复合操作 到 LiteRT Fused LSTM 操作转换。请参阅 LSTMCellSimple 转化代码。
转化合同:
Weight 和 projection 张量会被转置。
{input, recurrent} 到 {cell, input gate, forgot gate, output} gate}。
{bias} 到 {bias} 通过对偏差张量进行切片提取来提取。
投影是通过对转置投影张量进行切片来提取的。
类似的转化是针对 LayerNormalizedLSTMCellSimple。
LiteRT 转换基础架构的其余部分,包括 MLIR 凭证 并且最终导出到 LiteRT Flatbuffer 的 资源。
已知问题/限制
- 目前仅支持转换无状态 Keras LSTM(默认 行为)。有状态 Keras LSTM 转换是未来的工作。
- 仍然可以使用 使用底层无状态 Keras LSTM 层, 用户程序这种 TensorFlow 程序仍然可以转换为 LiteRT。
- 双向 LSTM 目前被建模为两个单向 SequenceLSTM 操作。它将替换为单个 双向序列 LSTM 操作。