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

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

算子融合

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

输入卷积层与归一化融合

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

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

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

x = x / 255

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

y = wx+b

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

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++写的一个伪代码的例子:

// 矩阵乘法 + 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=γvarYmean+βY=γvarWx+Bmean+βY=varγWx+varγ(Bmean)+β

常量折叠

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

更复杂的例子涉及到多个变量,比如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,就能提高效率。

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

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

死代码消除

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

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

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

总结

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

总结

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

11-15 19:39