TensorFlow Lite 转换器中的 TensorFlow 算子融合

文 / Ashwin Murthy,软件工程师,TensorFlow 团队

概览

效率和性能对于边缘部署至关重要。TensorFlow Lite 将更细化的一系列 TensorFlow 算子(本身由复合算子组成,如 LSTM)融合并优化单个可执行的 TensorFlow Lite 单元中,从而在效率和性能上达到理想效果。

此前,已经有诸多用户要求我们提供更细化的算子,以进一步提升性能。今天,我们对用户的需求做出了响应,为用户提供了融合指定算子的功能。

此外,这项新功能还支持 TensorFlow Keras LSTM 算子之间的无缝转换,这也是呼声最高的功能之一。更为方便的是,您现在还可以将用户定义的循环神经网络 (RNN) 转换插入 TensorFlow Lite!

让算子融合更加高效

如前文所述,TensorFlow 算子通常由多个更加细化的原始算子组成,例如 tf.add。这样的设计对于实现一定程度的复用性非常重要,并且可让用户根据现有单元自由组穿件算子。复合算子的一个例子是 tf.einsum。执行复合算子与执行组合中的每个算子的效果相同。

但是,如果要满足效率需求,我们通常会将一组更细化的算子计算“融合”到单个算子中。

融合算子的另一项用途是提供高阶接口,以定义量化等复杂转换,否则此类转换将无法实现,或难以在更细化的层面上完成。

TensorFlow Lite 中融合算子的具体示例包括各种 RNN 算子,如单向和双向序列 LSTM、卷积(conv2d、加偏置、ReLU)以及全连接(Matmul、加偏置、ReLU)等。

到目前为止,将 TensorFlow 的算子和 TensorFlow Lite 的算子进行融合,仍具有相当的 挑战性

开箱即用的 RNN 转换和复合算子支持

开箱即用的 RNN 转换

现在,我们支持将 Keras LSTMKeras 双向 LSTM 转换为复合 TensorFlow 算子。如要获取基于 RNN 的模型以利用 TensorFlow Lite 中的高效 LSTM 融合算子,这是最简单的方式。请参阅 此 Colab,了解如何通过 TensorFlow Lite 解释器进行端到端 Keras LSTM 到 TensorFlow Lite 的转换和执行。

此外,我们通过提供连接至基础架构转换的便捷接口,实现了到其他任意 TensorFlow RNN 实现的转换。您也可查看此功能与使用 lingvoLSTMCellSimpleLayerNormalizedLSTMCellSimple RNN 实现相关的数个示例。

如需获取更多信息,请查看我们的 RNN 转换文档

注:目前,我们致力于对 TensorFlow Lite 的 LSTM 算子添加量化支持。我们将在未来推出相关成果。

扩展至其他复合算子

我们扩展了 TensorFlow Lite 转换器,以便将其他复合 TensorFlow 算子转换为现有或自定义的 TensorFlow Lite 算子。

要实现 TensorFlow 算子融合至 TensorFlow Lite,需执行以下步骤:

  1. 将复合算子打包至 tf.function 中。在 TensorFlow 模型源代码中,使用 experimental_implements 函数注释标识复合算子并将其抽象为 tf.function。
  2. 编写转换代码。从概念上看,转换代码用已融合算子替换了此接口的复合实现。在 prepare-composite-functions 传递中,插入 转换代码
  3. 调用 TensorFlow Lite 转换器。使用 TFLiteConverter.from_saved_model API 转换为 TensorFlow Lite。

有关此基础架构的整体架构,请参阅 此处。有关代码示例的详细步骤,请参阅 此处。如需了解后台中的算子融合工作方式,请参阅 详细文档

反馈

请发送邮件至 [email protected],或使用 标签 “TFLiteConverter” 创建 GitHub issue

致谢

此项研究的顺利完成离不开 Renjie Liu 倾力支持,自项目成立以来,Liu 一直是主要协作者。同时,我们要感谢 Raziel Alvarez 的领导和指导。我们还要感谢 Jaesung Chung、Scott Zhu、Sean Silva、Mark Sandler、Andrew Selle、Qiao Liang 和 River Riddle 作出的重要贡献。最后,感谢 Sarah Sirajuddin、Jared Duke、Lawrence Chan、Tim Davis 和 TensorFlow Lite 团队,以及 Tatiana Shpeisman、Jacques Pienaar 和 Google MLIR 团队对这项研究的积极支持。