作为 AI 算法部署工程师，日常工作的其中一项就是将训练好的模型部署到设备上。根据业务需求的不同，需要保证推理的速度以及模型的稳定性。
在深度学习推理方面有多种提速方法，如模型剪枝量化与层算子融合等。算子融合是可以在深度学习中进行的一种典型优化，本文将配合 OpenPPL 高性能推理引擎来介绍应对动态模型的 Shape 算子融合策略。
OpenPPL 能够让人工智能应用高效可靠地运行在现有的 CPU、GPU 等计算平台上，为云端场景提供人工智能推理服务。OpenPPL 的自定义算子 PPL. ShapeOperation 可以有效将 Shape 与紧随其后的 Gather, Slice, Concat, Squeeze Unsqueeze, Add, Div 等算子进行融合，从而简化原始网络，提炼模型含义。
一、原始模型的问题
在 ONNX 官方定义中，Shape 算子输出的是输入 Tensor 的形状。Shape 的结果不参与核心的计算，但对整个推理过程至关重要。通常 Shape 算子会搭配 Gather, Slice, Add, Div, Concat 等算子组合使用，以实现数据定位、参数传递和矩阵变形等功能。
和其他算子组合使用的特性导致 Shape 相关操作十分臃肿，推理框架在处理这部分算子时发现两个棘手的问题：
- Shape 相关操作包含大量 kernel，虽然每个 kernel 都是简单的计算，但多次启动 kernel 和读写数据依然会对性能产生一定影响；
- Shape 算子得到的 Tensor 形状有时会经过连续的数值计算，之后作为参数传递给下一个算子，组合在一起的算子影响我们准确判断这部分算子的功能，从而影响我们后续的融合过程。
综上所述，我们需要自定义一种可以有效融合Shape算子并能简洁的表示其计算逻辑的新算子。
二、Shape 算子融合策略
2.1 等价替换 Shape 算子
自定义算子 PPL.ShapeOperation 内置一个独立的系数矩阵表示每个维度对应的系数，最后一维为常数项。以四维的 Shape 算子举例，等价替换原始的 Shape 算子，只需将对角线设置为 1，其他位置设置为 0，此时我们得到系数矩阵：
算子执行阶段时，PPL.ShapeOperation 首先提取输入 Tensor 的形状信息（以 NCHW 为例），生成输入向量：
之后，和系数矩阵的转置相乘得到算子输出结果：
在没有融合其他算子的情况下，其结果和原始 Shape 算子一致。
2.2 融合后续算子
可以融合的算子包括提取、拼接、数值计算三类。
提取用 Gather 和 Slice 算子通过指定维度的方式实现。PPL.ShapeOperation 会根据算子输入参数提取其中的行数（当输入参数为动态参数时不进行融合）。比如，Gather 算子取第二维（参数 index = 1），则原始的 M 矩阵会转换为：
拼接用 Concat 算子实现。当 Concat 需要拼接两个 Shape 的算子的结果时，只需要将两个系数矩阵进行拼接即可。比如，提取 Tensor 的第 2 维和第 3 维，模型会通过两个 Gather 算子得到两个维度的数据之后进行拼接，而在 PPL.ShapeOperation 中，初始矩阵 M 融合过程如下：
数值计算包含加减乘除，其中矩阵和常数的计算较为常见，以下图为例：
转换为数值公式为 (C+2)/2，则初始矩阵 M 融合过程如下：
融合后计算所需常数项均被保存在系数矩阵中，输入Shape与相乘的结果为上述四个算子的推理结果。
三、Shape 融合在 Shufflenet 中的应用
Shufflenet 网络需要将 Channel 平分为两份，之后将两份 Channel 交替排列，其中关于 Shape 的操作包含 Gather, Div, Concat。下图左侧为原始 Shufflenet 网络，右侧为融合之后 Shufflenet 网络结果，与其他 OpenPPL 融合策略共同作用下，融合共计 20 个算子。
融合后，PPL.ShapeOperation 算子维护两个稀疏矩阵，给第一个 Reshape 算子的系数矩阵为：
给第二个 Reshape 算子的系数矩阵为：
Shape 算子融合的第二个优势就是可以快速提取模型的功能，以上述两个系数矩阵为例：
- 第一个系数矩阵的功能为平分 C 轴。将第二维替换为 2，第三维替换为 1/2C。
- 第二个系数矩阵的功能为合并 C 轴。第二维数据为 -1，在其它轴大小固定的情况下，-1 会被通过以下公式进行替换为 C：
输入节点数已知输出节点数经过融合后，我们可以通过系数矩阵明确这几个算子在进行 ChannelShuffle 的操作，从而将其进一步融合为 ChannelShuffle 算子。
