Press "Enter" to skip to content

『算法理论学』深度学习推理加速方法之网络层与算子融合

 

任何事物都有连续性 –《极简主义》范式三:保持连续性的思维可以事半功倍

 

0. 引子

 

在深度学习推理方面有多种提速方法,如模型剪枝量化与层算子融合等。

 

网络层与算子融合是非常有效的方法,本文将配合 TensorRT 与 tflite 推理框架介绍下网络层与算子融合的原理与应用。

 

1. 融合理论

 

下面配合 TensorRT 介绍下网络层与算子融合的原理。

 

这是一个原始的 Inception Block,首先 input 后会有多个卷积,卷积完后有 BiasReLU ,结束后将结果 concat 到一起,得到下一个 input 。我们一起来看一下使用 TensorRT 后,这个原始的计算图会被优化成了什幺样子。

 

首先,在没有经过优化的时候 Inception Block 如 Figure1 所示:

 

 

第二步,对于网络结构进行垂直整合,即将目前主流神经网络的 conv、BN、Relu 三个层融合为了一个层,所谓 CBR,合并后就成了 Figure2 中的结构。

 

 

第三步,TensorRT 还可以对网络做水平组合,水平组合是指将输入为相同张量和执行相同操作的层融合一起。

 

 

最后,对于 concat 层,将 contact 层的输入直接送入下面的操作中,不用单独进行 concat 后在输入计算,相当于减少了一次传输吞吐,然后就获得了如 Figure4 所示的最终计算图。

 

 

通过上述的一些操作,网络图由 Figure1 简化为了 Figure4 的形式。通过融合的操作,使得

 

1. 网络层数减少,数据通道缩短

 

2. 相同结构合并,使得数据道路变宽

 

3. 更高效利用 GPU 资源

 

2.tflite 算子融合

 

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

 

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

 

让算子融合更加高效

 

 

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

tf.add
tensorflow.google.cn/api_docs/py…
tf.einsum
tensorflow.google.cn/api_docs/py…

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

 

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

 

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

 

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

 

挑战性

tensorflow.google.cn/lite/conver…

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

 

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

 

Keras LSTM

tensorflow.google.cn/api_docs/py…

Keras 双向 LSTM

tensorflow.google.cn/api_docs/py…

此 Colab

colab.sandbox.google.com/github/tens…

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

 

lingvo

github.com/tensorflow/…

LSTMCellSimple

github.com/tensorflow/…

LayerNormalizedLSTMCellSimple

github.com/tensorflow/…

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

 

文档

tensorflow.google.cn/lite/conver…

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

 

扩展至其他复合算子

 

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

 

TensorFlow Lite 转换器

tensorflow.google.cn/lite/conver…

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

 

 

    1. 将复合算子打包至 tf.function 中 。在 TensorFlow 模型源代码中,使用 experimental_implements 函数注释标识复合算子并将其抽象为 tf.function。

 

    1. 编写转换代码 。从概念上看,转换代码用已融合算子替换了此接口的复合实现。在 prepare-composite-functions 传递中,插入转换代码。

 

    1. 调用 TensorFlow Lite 转换器 。使用 TFLiteConverter.from_saved_model API 转换为 TensorFlow Lite。

 

experimental_implements
github.com/tensorflow/…
转换代码
github.com/tensorflow/…
TFLiteConverter.from_saved_model
tensorflow.google.cn/api_docs/py…

-1. 参考

 

-1.0: developer.nvidia.com/zh-cn/tenso…

 

-1.1: zhuanlan.zhihu.com/p/110934202

 

-1.2: mp.weixin.qq.com/s/KrlMxmO6X…

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注