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

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

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

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

深度学习性能优化之图优化_第1张图片

深度学习性能优化之图优化_第2张图片

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

LayerNorm算子合并

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

深度学习性能优化之图优化_第3张图片深度学习性能优化之图优化_第4张图片

模型中的不同的layernorm实现

深度学习性能优化之图优化_第5张图片

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

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

深度学习性能优化之图优化_第6张图片

深度学习性能优化之图优化_第7张图片

合并为ReduceSumSquare

同一个输出连接的相同计算合并为一个,特别是转换算子,如transpose, reshape等等。

深度学习性能优化之图优化_第8张图片

深度学习性能优化之图优化_第9张图片

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

删除没有必要的算子

如identity

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

深度学习性能优化之图优化_第10张图片

shape相关的常量折叠

在shape固定已知的情况下,shape可以提前计算得到替换为const,从而shape后续连接的算子可以进行常量折叠。

深度学习性能优化之图优化_第11张图片

transpose算子相关的优化

transpose算子优化的几种常见场景_Luchang-Li的博客-CSDN博客

算子顺序调节

调整顺序从而把div折叠到matmul和bias_add参数里面,第二个类似。这里可以考虑抽象一种比较通用的算法来处理更加广泛的情况。

深度学习性能优化之图优化_第12张图片深度学习性能优化之图优化_第13张图片

多路相同计算patten合并为batch计算,显著降低算子数量,提升算子计算密集性,如下图的场景;此外另一个场景是transformer模型中attention三个矩阵乘合可以并成为batch matmul,或者自己实现的特殊能够同时接受多个输入和bias的矩阵乘。

深度学习性能优化之图优化_第14张图片

多路并行的slice在特定情况下可以换成gather,并且后面所有的这些elemwise都可以合成一路计算。下图中这个稍微复杂点,slice可以替换为gahter,elemwise可以跟着合并,但gatherV2合并需要自定义算子batch gatherV2。 

深度学习性能优化之图优化_第15张图片

相邻的卷积/matmul+bias,中间没有非线性层的话可以基于卷积和矩阵乘的线性性直接合并

深度学习性能优化之图优化_第16张图片

算子替换

深度学习性能优化之图优化_第17张图片

只 slice一个axis的话可以换成gather从而简化实现方式,使用更高效的实现。

onnx和tensorflow pb模型图优化方法

onnx模型图优化/模型修改_Luchang-Li的博客-CSDN博客_onnx.checker.check_model

TensorFlow pb模型修改和优化_Luchang-Li的博客-CSDN博客_tensorflow模型优化

你可能感兴趣的:(深度学习编译器,深度学习,优化,算子,推理,性能)