部署神经网络时计算图的优化方法

部署神经网络时，各路框架基本都会把神经网络的计算建模为一个（有向无环的）计算图，之后再对这个计算图进行优化，包括硬件相关的优化和硬件无关的优化。本文介绍几种部署神经网络时计算图的优化方法，帮助读者在部署神经网络时理解部署工具都干了些什么。

算子融合

最关键的优化计算图的方式就是算子融合了，算子融合指的是将多个神经网络算子（例如卷积、池化、归一化等）组合在一起，以提高计算效率和性能。

输入卷积层与归一化融合

卷积神经网络中，输入的图像往往要做一个Normalization，比如ImageNet上训练的神经网络经常需要进行下面这个操作：

python 复制代码

std = [0.229, 0.224, 0.225]
mean = [0.485, 0.456, 0.406]
x = (x/255 - mean) / std

而YoloV5这样的模型则更简单，mean是0，std是1：

python 复制代码

x = x / 255

在第一层卷积时，我们会将卷积核在图像上进行滑动，用卷积核的参数乘以对应位置的像素（然后加上偏置），即

python 复制代码

y = wx+b

我们把Normalization的过程代入上面这个式子并化简：

python 复制代码

y = w(x/255-mean)/std + b
# 化简后
y = (w/255/std)x + (b-w*mean/255/std)

这个过程中，我们发现式子可以看作一个新的y=wx+b，新的w是w/255/std，新的b则是b-w*mean/255/std。所以，我们可以提前对卷积核参数进行变化，从而将两个算子合并为一个算子。

这是一个最简单的例子，其它的算子融合和上述操作的原理是类似的。

矩阵乘法操作 + 激活函数融合

神经网络里最常见的还有矩阵乘法操作 + 激活函数的组合，比如卷积层后面紧跟着一个ReLU，或者Transformer的FFN中Linear跟着一个GeLU等。激活函数可以在计算矩阵乘法的同时计算，下面是用c++写的一个伪代码的例子：

c++ 复制代码

// 矩阵乘法 + ReLU的融合
void MatMul(input, weight, bias, output, num_of_elements) {
    for (int i = 0; i < num_of_elements; i++){
        for (int j = 0; j < num_of_elements; j++){
            int accumulator = 0;
            for (int k = 0; k < num_of_elements; k++){
            	accumulator += input[i][k] * weight[k][j];
            }
            // relu
            output[i][j] = accumulator + bias[j] > 0 ? accumulator + bias[j]: 0;
        }
    }
}

矩阵乘法就是不断计算向量的点积，即一系列数字相乘后求和，而激活函数基本是对数字进行一个非线性的变换，所以非线性变换的过程可以求和得到结果之后马上进行，而不是进行完所有的矩阵乘法之后，再读一遍矩阵进行非线性的变换。

线性变换 + BatchNorm 融合

线性变换+BN也是一个常见的组合，比如YoloV5中就有很多的CBS（Conv+BN+SiLU）的组合（这里把Conv就是局部的线性变换）。

这种融合可以通过简化合并公式的方式进行：
Y = W x + B Y ′ = γ Y − m e a n v a r + β Y ′ = γ W x + B − m e a n v a r + β Y ′ = γ v a r W x + γ v a r ( B − m e a n ) + β Y=Wx+B \quad Y^\prime=\gamma\frac{Y-mean}{var}+\beta \\ Y^\prime = \gamma\frac{Wx+B-mean}{var}+\beta \\ Y^\prime = \frac{\gamma}{var}Wx + \frac{\gamma}{var}(B-mean) + \beta Y=Wx+BY′=γvarY−mean+βY′=γvarWx+B−mean+βY′=varγWx+varγ(B−mean)+β

常量折叠

这个是在代码编译领域也常用的优化方法，举个最简单的例子，当我们写python的时候，我们想让某个程序暂停两个小时，则一般会写time.sleep(2*60*60)，这个时候假如让程序在运行时再计算2*60*60就会略显繁琐，所以编译器会提前把2*60*60用7200代替。

更复杂的例子涉及到多个变量，比如a=5;b=a+3;c=b+a，这个时候优化的方式就是提前计算出b,c的值。

公共子表达式消除

公共表达式是传统编程语言编译器常用的优化的一种，在程序中计算表达式时，有时会出现公共的子表达式，重复计算这些子表达式会增加计算开销。

比如下面这个例子：

temp = b * c
a = b * c + g
d = b * c + e

计算a和d的时候会重复计算b*c，假如我们只计算一次temp = b * c，然后计算a=temp+g, d=temp+e，就能提高效率。

到了神经网络领域，可能会出现下图左边这个情况，此时需要合并成右边这种情况，从而简化计算图。

死代码消除

这也是传统编程语言编译器中的一种优化方式，比如下面这个代码，return之后的代码是unreachable的，所以编译后应该完全消除这部分。

python 复制代码

def test(flag):
    print("Flag is False.")
    return
    print("This code is unreachable.")

到了神经网络上，在计算图上其实能比较直观地发现没有用的节点或者不可达的节点。比如有一个节点孤立在计算图外面；或者某个节点有输入没输出且不是输出节点；或者某个节点有输出没输入且不是输入节点。

总结

神经网络编译器和传统编程语言编译器非常相似，其许多优化技术都是从编程语言编译器中沿用而来，但是神经网络编译器也有它的特点，有新的例如算子融合的优化方法可以用。这些优化方式能够对神经网络的部署起到关键的作用。

部署神经网络时计算图的优化方法