秋刀鱼在做梦

这个屌丝很懒，什么也没留下！

热门标签

深度学习性能优化之图优化

作者：秋刀鱼在做梦 | 2024-08-03 06:41:50

踩

图优化

这里总结了作者对一些实际深度学习模型优化实施的部分图优化，有一些现有的深度学习框架已经有了，有些是作者的独特发现，现有深度学习框架还没有。后续有新的发现将持续更新，也欢迎读者贡献有较为广泛应用场景的独特优化。

基本概念：深度学习模型是由算子连接构成的有向无环图(DAG)，我们可以对这个图的子图替换为一个计算等价的另一个子图，从而对这个子图实施更加高效的计算。常规的图优化是用框架已有的算子来组合新的子图替换旧的子图，而算子融合是图优化的一种特殊情况，它采用人工或自动生成的新的算子来替换旧的子图。

为什么会存在图优化的机会？有很大一部分是算法开发人员对底层部署并不是很熟悉，算法层面可以非常灵活地用不同的方法来表达同一种计算，但是不同计算方式在推理引擎和硬件底层可能具有显著不同的性能（当然不同的硬件往往需要不同的优化方法）。另一方面跟引擎有关，例如LayerNorm本身为一个算子，但是ONNX导出会拆分为很多细小的算子，则必须用图优化将其合并回去。又比如ONNX不支持NHWC等数据格式，但是硬件部署往往又需要转换到硬件优化的NHWC或其他数据格式。另外，根据引擎和芯片的特殊性，往往也需要对模型进行一些针对性修改优化，例如一些引擎不支持超过5D的transpose计算。

如何做图优化？可以在多个阶段进行图优化操作。第一个阶段是直接对算法代码进行修改，从而导出的模型直接具有优化后的特性。对于一些算法依赖，通用性并不是很强，或者底层实施图优化困难的优化可以考虑该方法。第二个阶段是对导出的模型进行优化，例如对导出的ONNX模型进行优化。第三个阶段是在深度学习引擎内部运行时进行优化，可以结合硬件特点和输入shape对应的性能信息进行更加针对性和极致的优化。

ONNX模型进行图优化的编程方法参考：

onnx模型图优化/模型修改

OneHot+MatMul to Gather

下图bert模型里面这个pattern实际上能换成一个gather计算，可以减少算子数量、种类，以及显著简化计算方式。

有时候多个gather/slice可以替换为split，减少算子数量。

SpaceToDepth与DepthToSpace算子

图中reshape+transpose(perm=[0, 1, 3, 5, 2, 4])计算等价于SpaceToDepth，由于transpose场景特别多，这里替换后可以在计算上进行更加针对性优化，此外，这个优化使得做NCHW到NHWC等格式转换优化也更加容易。

跟上面相似，这里transpose(perm=[0, 1, 4, 2, 5, 3])等价于DepthToSpace计算。

当然6D的transpose转换为SpaceToDepth和DepthToSpace可能不止一种perm，可以参考Onnx的这两个算子定义还提供了其他的perm可能性。例如还有这种channel维度放在后面的：

这个pattern其实也是在做一种space2depth

但是标准的space2depth是把NCHW解释为NCH1H0W1W0，再transpose为NH0W0CH1W1，H0W0为 space2depth的block size。但是上面的pattern transpose类型可能跟标准的space2depth不一样，例如排布可能是NW0H0CH1W1，但是仍然可以直接替换为reshape+transpose+reshape。

下图对应于pytorch的roll计算：

x = torch.roll(x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))

transpose+reshape+transpose合并

更加复杂的情况，reshape算子同时存在拆分和合并元素：

可以把上述任意多个transpose+reshape+transpose合并为一个reshape+transpose+reshape的pattern。

一种比较简单的融合任何transpose+reshape+transpose的算法为：

把输入shape拆分为最小的素数单元，例如[4, 9]拆分为[a0,a1,a2,a3]。我们需要用一个符号shape来表示拆分结果，每个符号shape对应一个实际的拆分的素数结果。

这里存在的一个问题是，比如6可以拆分为[2,3]也可以拆分为[3,2]，有的拆分方式可能导致后续无法进行处理。比如6拆分为[2,3]，reshape为[3,2]时，无法处理，但是把6拆分为[3, 2]可以。可以考虑根据不同拆分方式进行尝试。

再根据输入shape把这些小单元合并为实际的shape组,例如[4,9]为[[a0,a1], [a2,a3]]

根据这个小的shape分组进行推导输出的shape格式，例如为[[a2,a3], [a0,a1]]

把所有小的shape合并成一个整体并且根据符号shape标号的连续性进行合并分组，得到[a2,a3], [a0,a1]

根据这个信息很容易得到最终合并的reshape和transpose信息。

这些算子中间带有一些种类的elemwise一样可以进行合并。

合并效果演示：

不是所有的场景都能融合成一个transpose，比如这个场景，可以考虑只合并其中一部分。

替换为clip算子：

split与bias add交换顺序

该场景可见CLIP text encoder

这个mul算子可以放到下面的reshape和 transpose间的任何位置，可以放到合适的位置从而跟其他算子更好的融合。

矩阵乘+BN融合

跟卷积+BN融合类似Conv2D + batch normalization (BN) 融合_Luchang-Li的博客-CSDN博客_bn融合

反卷积+BN也应该能进行融合

合并相邻的Conv2D或MatMul

这两个算子中间不能有非线性算子等计算，可以直接把两个相邻的Conv2D和矩阵乘合并为一个。

特殊的1x1 depthwise卷积替换为elemwise

-》

K=1的矩阵乘，直接替换为Mul算子。

MatMul与Add, Mul向量计算融合

图中MatMul和这后面的三个Add和Mul算子和融合为一个MatMul+add，推广开来实际上MatMul后面可以融合任意多个Add和Mul算子，但是有一些前提，比如Add和Mul第二个输入需要是常量，并且维度是1维或二维等。

例如融合MatMul+bias再加一个Add或Mul向量：

(in*B+bias0) + bias1 = in*B + (bias0+bias1)

(in*B+bias0) * sclae1 = in*(B*sclae1) + (bias0*sclae1)

因此可以一直往后融合下去（第一个MatMul没有bias可以认为其包含一个内容为0的bias）。

这个融合在性能上肯定是提升的，但是在精度上有可能会产生牺牲，特别是融合了mul算子。因为原来的矩阵乘weight和bias以及scale通常都是比较小的小于1的数值，把scale融合到weight和bias里面后，会导致weight和bias数值进一步降低，可能导致精度下降。

可以交换mul和add，把add内容除以mul内部标量值，从而使得add可以以bias融合到matmul中，同时后续的mul也可以融合到matmul中。但是FP16计算可能精度损失较大。

DepthwiseConv+Add

计算为x*w1+bias1+x,可能可以转换为x*w1+x*w2+bias1，这里*为卷积。

其中w2为delta函数，也就是中间为1，其余值为0的filter。因此最终这个计算可以转换为

x*(w1+w2)+bias1

(X+b1)*w2 = X*w2+b1*w2

可以把第一个Add移动到MatMul下面来，然后下面两个add可以合并成一个。

LayerNorm算子合并

一些深度学习框架的Layernorm是多个小算子组成的，把上面的这个pattern识别进行替换为一个手写的融合layernorm算子，融合后的这个算子性能可以提升好几倍。onnx最新的opset 17才增加了这个融合的算子定义。

模型中的不同的layernorm实现

这是另一种layernorm的情况，这主要可能是代码开发问题，导致reduce后面出现两个sub算子，但是这两个算子计算是一模一样的，把这两个算子合并后就跟前面两种情况一样了。

个别模型里面pb转onnx时layernorm的reduce被替换为GlobalAvgPool，在shape长度为3时可以直接替换为reducemean从而避免上面layernorm匹配不成功问题。

归一化函数合并

这一堆算子其实由F.normalize(x, p=2.0, dim=-1)导出，可以进行针对性合并。第二个图计算跟第一个图类似，但是为不同代码编写方式，这种最好是修改模型pytorch代码改成同一种norm函数，降低模型图优化代码开发。

GroupNormalization

reshape+InstanceNorm+reshape

这里实际上是个GroupNormalization算子，广泛用于stable diffusion，但是，GroupNormalization在onnx opset 18才首次出现，这可能是这里采用instanceNorm替代的原因。

Global Response Normalization (GRN)

https://github.com/open-mmlab/mmpretrain/blob/main/mmpretrain/models/utils/norm.py#L10

这里三个Slice实际上可以替换为一个Split，当然也可以看哪种实现方法更好

合并为ReduceSumSquare

下图这里可以改成reshape+transpose计算更高效，当然最好是构建模型代码层面上去修改。这里深刻体现了同一种计算范式，不同的代码开发者实现的多样性，但是同种计算不同构图方式的计算效率可能是不太一样。

合并相同计算

特别是转换算子，如transpose, reshape等等。

如何更加智能高效的识别这些相同的计算pattern也是一个值得研究的问题，而不是每次都要人为去判断。

删除没有必要的算子

如identity，expand

下图中的expand其实没有必要，因为where作为一个ternary elemwise自带broadcast能力，而expand导致需要额外读取内存并写回一个很大的tensor导致时间浪费。可以进行如下通用性判断：expand后面接elemwise算子并且这个elemwise的其他输入shape与expand的输入shape能broadcast到elemwise输出的shape时，这个expand可以删除。

可以直接变reshape，而reshape一般不需要实际计算。

split concat tensor合并

上图这种情况也可能是代码开发导致的，split了很多个小tensor，然后又concat在一起，而concat中来自split的输入是连续的，因此可以把中间split的tensor合并在一起。也就是说上面图中split了24个tensor，实际上只需要split成2个tensor。第一个concat可以直接去掉。