炼丹中缺失的一课

阅读数: 次 2023-06-10

当我大学入学时，已是2020年。那时候已经有了许多深度学习框架，Tensorflow, Pytorch, Caffe等。同时计算机技术的发展，甚至让我可以在我的笔记本玩《方舟：生存进化》。所以当我误入歧途开始炼丹的时候，我们好像并不用那么仔细的考虑，炼丹的复杂度。这一系列隐藏在封装好的代码和底层的事情，我们就管它们叫《炼丹中缺失的一课》吧。

当我们想要描述一个深度学习模型的复杂度时，我们一般从2个角度来切入，参数量和计算量。针对这两个方面，一般有如下4个指标。

metrics	description
Params	模型的参数量
FLOPs	FLoating point OPerations，浮点运算次数。（注：不要将其与FLOPS混淆，后者是每秒浮点运算次数的意思，衡量的是硬件的性能。）
MAC	Memory Access Cost，内存访问成本
MACC	Multiply-ACCumulate operations，乘加操作次数

在深度学习中，我们大部分时候运行的操作都是卷积，矩阵乘法。实际上这也是最为耗时的环节。无论使用哪种运算的策略，都逃不开乘法和加法。例如一次3×3的卷积核的卷积，需要进行9次乘法和9次加法，所以一般1MACC≈2FLOPs。由于现在许多的硬件都将乘加运算当成一个单独的指令，所以我们可以只关注FLOPs。

Params & FLOPs

我们下面以一个朴素的卷积神经网络为情景，在2023年，一个CNN它应该具有卷积层，全连接层，BN层，激活函数层。特别地，我们会讨论卷积层中的分组卷积和深度可分离卷积。我们先考虑Params：

layers	prior	Params
Conv	输入通道数$C_i$ 输出通道数$C_o$ 卷积核尺寸$K$	$\underset{\boldsymbol{w}}{\underbrace{C_i\times K^2\times C_o}}+\underset{\boldsymbol{b}}{\underbrace{C_o}}$
GroupConv	输入通道数$C_i$ 输出通道数$C_o$ 卷积核尺寸$K$ 分组数$g$	$\underset{\boldsymbol{w}}{\underbrace{\frac{C_i}{g}\times K^2\times \frac{C_o}{g}\times g}}+\underset{\boldsymbol{b}}{\underbrace{C_o}}\\=\underset{\boldsymbol{w}}{\underbrace{\frac{1}{g}\times C_i\times K^2\times C_o}}+\underset{\boldsymbol{b}}{\underbrace{C_o}}$
Depth-wise Conv	输入通道数$C_i$ 输出通道数$C_o$ 卷积核尺寸$K$	$\underset{\boldsymbol{w}}{\underbrace{C_i\times K^2+C_i\times C_o}}+\underset{\boldsymbol{b}}{\underbrace{C_o}}\\=\underset{\boldsymbol{w}}{\underbrace{\left( \frac{1}{C_o}+\frac{1}{K^2} \right) \times \left( C_i\times K^2\times C_o \right) }}+\underset{\boldsymbol{b}}{\underbrace{C_o}}$
FC	输入神经元数量$N_i$ 输出神经元数量$N_o$	$\underset{\boldsymbol{w}}{\underbrace{N_i\times N_o}}+\underset{\boldsymbol{b}}{\underbrace{N_o}}$
BN	输入通道数$C_i$	$2\times C_i$（$\gamma, \beta$）
Activation	N/A	N/A

默认情况下，网络中的每个参数是32位浮点数的。1个fp32对应4个字节，所以一般模型大小是参数量大小的四倍。这就引出了通过量化(quantization)来压缩模型的操作，将fp32转化为int8，可以让模型大小直接缩小4倍，且大多数时候都不会特别伤害性能，在大部分的嵌入式部署时，这都是绝对值得做的一项优化。

下面我们考虑FLOPs的计算，我们沿用上面的符号，对于卷积的情形，我们记输入图像大小为$(H_i,W_i)$，输出图像大小为$(H_o,W_o)$。我们应该非常熟悉卷积的过程，对于$(C_i,H_i,W_i)$的数据，我们通过$C_o$个滤波器组，每组滤波器有$C_i$个卷积模板，来让数据转换为$(C_o,H_o,W_o)$。我们知道，对于一个寻常的卷积过程：

$H_o=\lfloor \frac{H_i-K}{S} \rfloor +1 \\ W_o=\lfloor \frac{W_i-K}{S} \rfloor +1$

所以，每个滤波器组里的$C_i$个卷积核，会与输入数据的$C_i$个通道分别相乘。所以一次这样的操作需要进行$C_i\times K^2$次乘法，$C_i\times K^2-1$次加法。这样的滤波器组共有$C_o$个，且这样的一次操作只是计算出特征图上的一个点，一共有$H_o\times W_o$个点需要计算。同时，如果考虑偏置，偏置的加法运算显然进行了$C_o\times H_o \times W_o$次。所以最终的FLOPs为：

$=\left( C_i\times K^2+C_i\times K^2-1 \right) \times C_o\times H_o\times W_o+C_o\times H_o\times W_o \\ =2\times C_i\times K^2\times C_o\times H_o\times W_o$

观察没有合并前的式子，如果我们将乘法和加法合并为同一种操作，那么MACC即为$C_i\times K^2\times C_o\times H_o\times W_o$。

类似的，对于分组卷积的情况，式子可以修改为：

$=\left[ \left( \frac{C_i}{g}\times K^2+\frac{C_i}{g}\times K^2-1 \right) \times \frac{C_o}{g} \right] \times g\times H_o\times W_o+C_o\times H_o\times W_o \\ =\left( \frac{C_i}{g}\times K^2+\frac{C_i}{g}\times K^2-1 \right) \times C_o\times H_o\times W_o+C_o\times H_o\times W_o \\ =\frac{1}{g}\times \left( 2\times C_i\times K^2\times C_o\times H_o\times W_o \right)$

在计算深度可分离卷积时，我们知道，这个计算可以被分成两个部分，depth-wise和point-wise。

对于$C_i$通道的输入的数据，每个通道仅由一个卷积核负责。所以此时的一次卷积会进行$C_i\times K^2$次乘法，$C_i \times K^2 -1$次加法。然后这$C_i$个中间结果会用1×1的卷积核将通道修改为$C_o$。由于1×1的卷积并不会改变特征图的尺寸，所以我们知道depth-wise阶段的特征图尺寸即$(H_o,W_o)$。所以depth-wise阶段的FLOPs即$=2\times C_i\times K^2\times H_o\times W_o$。在point-wise阶段，我们可以认为是$K=1$的普通卷积，所以FLOPs为$=2\times C_i\times C_o\times H_o\times W_o$。所以最终的FLOPs是：

$=2\times C_i\times K^2\times H_o\times W_o+2\times C_i\times C_o\times H_o\times W_o \\ =\left( \frac{1}{C_o}+\frac{1}{K^2} \right) \times \left( 2\times C_i\times K^2\times C_o\times H_o\times W_o \right)$

BN层的FLOPs在训练和测试时是不同的，我们这里就只考虑在测试的情形了。在测试时：

$y=\frac{\gamma}{\sqrt{\mathrm{Var}\left[ x \right] +\epsilon}}\cdot x+\left( \beta -\frac{\gamma \mathrm{E}\left[ x \right]}{\sqrt{\mathrm{Var}\left[ x \right] +\epsilon}} \right)$

这里的系数此时的都是已知的，所以BN层的FLOPs易得即$2\times H_i\times W_i\times C_i$。当然，训练时就比较复杂了，因为训练时涉及到样本均值和方差的滑动平均估计，但我们也并不关心训练时增多的那点FLOPs。另外，由于在测试时，上面的这些值都是已知的，所以在真实部署时，可以直接把BN的系数融合进前一层卷积中。这个操作也叫“吸BN”。

FC层的FLOPs也是容易计算的，沿用上面的定义，每个输出的神经元都需要由$y_i=\sum_{j=1}^{N_i}{w_jx_j}+b_i$计算得到，一次需要$2\times N_i$，一共有$N_o$个输出，所以最终是$2\times N_i \times N_o$。实际上，我们可以将全连接看作一个全局卷积：

$=2\times C_i\times K^2\times C_o\times H_o\times W_o \\ H_o=1 \\ W_o=1 \\ C_o=N_o \\ K^2=H_i\times W_i \\ C_i\times H_i\times W_i=N_i \\ =2\times N_i\times N_o$

激活函数的FLOPs是很直接的element-wise操作，在卷积的框架下，即为$C_i\times H_i \times W_i$。在全连接的框架下，那直接为$N_i$。

推导刚才的式子可以作为一种课堂练习，但本身式子的结果其实意义不大，因为我们在评估FLOPs和Params时往往面对的网络都很复杂，有复杂的跳连等等，真正评估时我们往往都是用一些工具或装饰器。所以上述推导是为了给出如下的几个很重要的见解：

FLOPs和输入图片的尺寸有关，所以输入小分辨率的图片往往可以加快速度。
分组卷积和深度可分离卷积通过分组的思想，减轻了普通卷积的耦合度，降低了参数量和FLOPs。（但我们后面将会看到，这样并不意味着运行起来更快。）深度可分离卷积的第二阶段，使用point-wise实现了多通道之间的信息交流，至于分组卷积，可以使用shuffle的方式，这即是ShuffleNet的做法。
对于一个朴素的基于卷积的神经网络，卷积和全连接是最为耗时的操作。激活函数，BN层并不会有很大的开销。（如今人们都会使用Global Average Pooling了，所以全连接层的参数量和FLOPs的开销反而没那么大了。）
跳连，逐元素相加/相乘这样的操作，可能不会带来参数量的增加，但会增加FLOPs。

MAC

现在，需要更细致的考虑“运行一个CNN”这件事情。我们运行一个CNN都必须依赖于具体的硬件，如CPU, GPU等。硬件有两个主要指标，一个是刚才的FLOPS，我们记作$\pi$；另一个是带宽，我们记作$\beta$。带宽指的是硬件一次每秒最多能搬运多少数据(Byte/s)，即内存交换量。所以我们可以定义一个比例系数—计算强度$I$。所以对于硬件平台来说：

$I_{\max}=\frac{\pi}{\beta}$

它的单位是FLOPs/Byte，描述了单位内存交换可以用来进行多少次计算。那么对于一个深度学习模型，我们可以按照上述定义给一个模型的计算强度：

$I=\frac{\mathrm{FLOPs}}{\mathrm{MAC}}$

我们考虑用CPU进行运算，此时输入的数据，网络权重本身都被读入到内存中了。由简单的计组的知识，我们知道，加法和乘法都只能通过算术逻辑单元ALU实现。所以这里就涉及到将权重和数据读入CPU，计算完后再写回内存了。由于内存相对于CPU的速度较慢，所以读入和写回操作会成为计算的瓶颈。这也导致，不能只通过Params和FLOPs来解释一个CNN的效率。这最早是在ECCV 2018的ShuffleNetV2中被提出的。接下来我们会沿申ShuffleNetV2中的4个guidelines，来扩展一些内容。

输入输出通道数相同时，内存访问成本最小。

原文以这样的一个toy model为例，考察1×1卷积，假定输入输出通道数为$C_{in},C_{out}$，特征图尺寸为$h,w$。那么CPU需要先读入输入$hwC_{in}$，权重$C_{in}C_{out}$，最后写回输出$hwC_{out}$。所以内存访问成本理论上是：

$MAC=hw\left( C_{in}+C_{out} \right) +C_{in}C_{out}$

根据均值不等式，容易得$C_{in}=C_{out}$时，MAC最小。虽然这是理论上的，但原文的实验也证明了这是正确的。

过多的分组卷积会增大内存访问成本。

这个结论的前提，是在给定计算量FLOPs时，仍考察上面的那个过程，FLOPs为$B=hwC_{in}C_{out}/g$，所以此时的MAC：

$MAC=hw\left( C_{in}+C_{out} \right) +\frac{C_{in}C_{out}}{g} \\ =hwC_{in}+\frac{Bg}{C_{in}}+\frac{B}{hw}$

但是这个解释其实很奇怪，这是在给定FLOPs的情况下。但大部分时候的场合都是，比如在应用nn.Conv2d()，我将groups设成8。这时候理所应当是快的。历史上人们发现比标准卷积更慢，实际上是底层算子没有优化到跟标准卷积一样好的原因。

网络碎片化的结构会减少并行度

例如GoogleNet的inception结构，和一些NAS出来的结构。这些多分支的结构由于涉及同步的问题，所以会降低速度，这是很好理解的。这点也可以在PyTorch中得到验证。

element-wise操作不能忽视

这些逐元素的操作，比如激活函数激活，跳连相加，计算偏置，以及以后的channel-attention。都是FLOPs比较小，但MAC相对较大。它们对速度的影响不能忽视。

结合这4个guidelines，和前面的前置知识。我们可以将部署一个神经网络模型，看作是从指标（metrics），大小（size），速度（speed）三个做trade-off。具体来说，如果我们的模型只是为了刷SOTA，在标准计算机，服务器上进行推理。那其实我们可以专心用各种技巧，比如疯狂跳连（姿态估计任务中），以及构造各种注意力。如果我们的模型是部署在移动端，移动端虽然算力不是那么有限，但内存还是充足的。使用跳连，内存上是可以满足的。所以只要在可以容忍的延迟之内，就可以成功部署。而在一些更极限的场所，比如微处理器，这时候内存是死线，所以最好就只设计一个单路的网络，直接进行推理。一个开源库TinyEngine，对深度可分离卷积进行了调整。让其只用一个原地的缓存，即用时间换空间。这个过程主要是一个“trial and error”的过程……

End

如果到实际部署上，其实还有两步，剪枝和量化。这里就先不提了，哪天有空再写吧。