TensorFlow RNN 转换为 LiteRT

概览

LiteRT 支持将 TensorFlow RNN 模型转换为 LiteRT 模型 Fused LSTM 操作。融合操作旨在最大限度地提高 底层内核实现,并提供更高级别的 API, 接口来定义量化化等复杂的转换。

由于 TensorFlow 中有许多 RNN API 的变体,我们的方法 双层:

  1. 提供对标准 TensorFlow RNN API 的原生支持,例如 Keras LSTM。 此为推荐选项。
  2. 提供连接 转换基础架构的 用户定义的 RNN 实现,用于插入并转换为 LiteRT。我们在这里提供了几个现成的示例, 使用 lingvo 的 LSTMCellSimpleLayerNormalizedLSTMCellSimple 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 到 LiteRT 的转换。对于 请参阅 Keras LSTM 接口 和转化逻辑 此处

还有一点也很重要,那就是强调 LiteRT 的 LSTM 合同 添加到 Keras 运算定义:

  1. input 张量的维度 0 是批次大小。
  2. recurrent_weight 张量的维度 0 是 输出。
  3. weightrecurrent_kernel 张量经过转置。
  4. 转置权重、转置 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,如下所示 如下:

  1. 使用 tf_implements 将 LSTMCellSimple 的所有使用过程封装在 tf.function 中 带有该标记的注释(例如 lingvo.LSTMCellSimple 应为 注释名称)。确保生成的 tf.function 与转化代码中预期函数的接口相符。这个 是添加了注解的模型作者与 转化代码。
  2. 扩展 prepare-composite-functions 传递以插入自定义复合操作 到 LiteRT Fused LSTM 操作转换。请参阅 LSTMCellSimple 转化代码。

    转化合同:

  3. Weightprojection 张量会被转置。

  4. {input, recurrent}{cell, input gate, forgot gate, output} gate}

  5. {bias}{bias} 通过对偏差张量进行切片提取来提取。

  6. 投影是通过对转置投影张量进行切片来提取的。

  7. 类似的转化是针对 LayerNormalizedLSTMCellSimple

  8. LiteRT 转换基础架构的其余部分,包括 MLIR 凭证 并且最终导出到 LiteRT Flatbuffer 的 资源。

已知问题/限制

  1. 目前仅支持转换无状态 Keras LSTM(默认 行为)。有状态 Keras LSTM 转换是未来的工作。
  2. 仍然可以使用 使用底层无状态 Keras LSTM 层, 用户程序这种 TensorFlow 程序仍然可以转换为 LiteRT。
  3. 双向 LSTM 目前被建模为两个单向 SequenceLSTM 操作。它将替换为单个 双向序列 LSTM 操作。