文 / 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 LSTM 和 Keras 双向 LSTM 转换为复合 TensorFlow 算子。如要获取基于 RNN 的模型以利用 TensorFlow Lite 中的高效 LSTM 融合算子,这是最简单的方式。请参阅 此 Colab,了解如何通过 TensorFlow Lite 解释器进行端到端 Keras LSTM 到 TensorFlow Lite 的转换和执行。
此外,我们通过提供连接至基础架构转换的便捷接口,实现了到其他任意 TensorFlow RNN 实现的转换。您也可查看此功能与使用 lingvo 的 LSTMCellSimple 和 LayerNormalizedLSTMCellSimple RNN 实现相关的数个示例。
如需获取更多信息,请查看我们的 RNN 转换文档。
注:目前,我们致力于对 TensorFlow Lite 的 LSTM 算子添加量化支持。我们将在未来推出相关成果。
扩展至其他复合算子
我们扩展了 TensorFlow Lite 转换器,以便将其他复合 TensorFlow 算子转换为现有或自定义的 TensorFlow Lite 算子。
要实现 TensorFlow 算子融合至 TensorFlow Lite,需执行以下步骤:
- 将复合算子打包至 tf.function 中。在 TensorFlow 模型源代码中,使用 experimental_implements 函数注释标识复合算子并将其抽象为 tf.function。
- 编写转换代码。从概念上看,转换代码用已融合算子替换了此接口的复合实现。在 prepare-composite-functions 传递中,插入 转换代码。
- 调用 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 团队对这项研究的积极支持。