为什么要进行模型压缩?
本文主要从算法(而非硬件)的角度出发,介绍一些常用的模型压缩技术。
深度学习网络通常是 over-parameterized ,参数是高度冗余的,并且不同部分的参数对性能的贡献差异也非常大。顾名思义,网络剪枝就是将模型中一些不必要的参数修剪掉。
网络剪枝的流程如下图所示:
先训练一个大的、准确的网络;
评估该网络参数中每个参数/每个神经元的 “重要性”;
如何度量参数/神经元的 “重要性” 呢?最直觉的想法:参数的重要性可以由它绝对值的大小来表征,绝对值较大的参数对网络整体的影响较大,比较重要;神经元的重要性可以通过记录对于多个输入,它输出为零的次数来表征,经常输出为零,说明它不太重要。
移除掉不重要的参数/神经元;
微调剪枝后的网络;
通常,在移除掉一些神经元之后,网络的性能会有所下降。这时,我们对剪枝过的网络进行微调,使得它的准确率回升一些
重复 3,4 步骤直到模型的大小和精度都满足我们的需求。
如果我们要移除参数,这个过程如下图所示。这样做得到的会是不规则的,这样会带来两个问题:一是不太好实现,想想我们在深度学习框架如 Pytorch 中定义网络层时都是直接指定神经元的个数,框架会给出一个标准的层。而如果要指定某些神经元的某些参数不存在,是比较麻烦的;二是不利于 GPU 加速,现代神经网络大都是通过 GPU 进行并行加速的,不管是全连接还是卷积,其底层实现都是进行矩阵乘法,而如果指定某些参数不存在,对于 GPU 加速也是很麻烦的。
想要实现移除参数,一个可行且直接的做法是:不是真的移除参数,而是将该参数的权重置为 0。这样做就能规避上面提到的两点问题,达到与移除指定参数同样的效果。但是,想想看,我们现在是在做模型剪枝,如果并没有实际移除参数,那网络的参数量和计算量都根本没有减小啊,这就没有意义了。
因此,实际中,通常采用的是移除神经元的方式。
如果通过移除神经元来实现模型剪枝,模型在剪枝之后还是规则的,这就没有上面移除参数方式中的难以实现和难以加速的两点问题。
当有一个准确度较高的大模型,通过剪枝的方式将它的参数量和计算量减小,而且掉点不多,这是模型剪枝做的事情。但是,为什么不直接训练一个比较小、也比较准确的模型呢?
实际上,大模型比较容易训练,并且也可以在剪枝之后保持准确率。而小模型很难通过直接训练得到较高的准确率。
The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks 提出了大乐透假说,给出了大模型训练、小模型训练与剪枝的一种解释。
文章认为,训练一个神经网络,某一组结构的某一组初始化参数是可以训练成功的,而其他结构、初始化参数,很可能就是无法训练成功的。就像买彩票一样,如果一个神经网络中恰好有这一组结构、初始化参数,那么这个网络就是可以训练成功的。而一个神经网络可以看作是很多子网络的结合,在这些子网络中,只要有一组训练成功了,那么整个大的网络就训练成功了。
大模型中子网络更多,相当于买彩票的注数更多,自然更可能成功。大模型成功之后,就已经证实了该模型中的一个子网络是可以成功的,因此如果将大模型中的其他子网络剪枝掉,就可以得到规模又小、准确率又高的网络。而如果直接训练小模型,相当于买彩票的注数较少,因此很难直接训练出一个性能很好的小网络。
大乐透假说是如何通过实验证明的呢?
下图中,先有一个随机初始化的大模型(最左侧),经过训练之后,得到高准确率的大模型。对训练好的大模型进行剪枝,得到的是大模型的一个子网络(小模型),并且这个小模型性能也比较好。然后对于这个小模型,如果完全重新随机一组初始化参数,对小模型从头训练,是无法得到好的性能的。但是,如果使用原始大模型的随机初始化参数对小模型进行初始化,再进行训练,就可以得到好的性能。
这说明,大模型通过训练,找出了一个可行的初始化参数和子网络,这就是 “中奖” 的那注彩票,验证了大乐透假说。
这里再列几篇与大乐透假说同期或之后的研究工作以及它们的主要结论,有兴趣读一下原文:
Deconstructing Lottery Tickets: Zeros, Signs, and the Supermask
Weight Agnostic Neural Networks
Rethinking the Value of Network Pruning
知识蒸馏是指先训练一个大的教师模型,然后用教师模型训练学生模型,将教师模型的输出 logits(知识)蒸馏给学生模型。这样做往往能够比直接从头训练一个小的学生模型的效果更好,原因已经在模型剪枝部分讨论过了。
知识蒸馏的常见流程如图所示。先用数据集中的样本及 one-hot 标签正常训练一个大的教师模型(Teacher Net),训练完成之后,固定教师模型的参数,将它的输出 logits 作为标签,训练一个相对较小的学生网络(Student Net)。
相比于 one-hot 标签,教师模型的输出能够提供更多、更丰富的信息量,比如不同类别之间的相近程度。虽然我们知道分类模型是要让同类的输出尽量相近、不同类的输出尽量远离。但在实际中,我们数据集中的各个类别之间的相近关系也有所不同,比如数字1和数字7就比较接近,和数字8就比较远离。然后,“硬”的 one-hot 标签并不能给出这种不同类间的关系。但是教师模型的输出却可以。这种更丰富的、更有意义的标签或许能够帮助学生模型学的不那么吃力:对这张图片不需要给出 数字1 是1,其他都是0的预测,给出 1: 0.7, 7: 0.2, 9: 0.1 … 的预测就很好了。
Distilling the Knowledge in a Neural Network
在知识蒸馏中所使用的 softmax 与常规 softmax 有所不同,在于对每个标签的值除以了一个温度系数 temperature T T T,这可以使得标签由陡峭变得平滑。如下图所示,左侧是常规 softmax,右侧是经过标签平滑的 softmax 。
为什么要让标签平滑呢? y i y_i yi 是教师模型的输出 logits,如果直接经过常规 softmax 处理,得到 y i ′ y'_i yi′ ,各类之间的差异会变的陡峭,这样的话其实就跟原本的 one-hot 标签差不多了,就失去了做知识蒸馏的意义。因此,通过除以一个温度系数 T T T,使得标签变得平滑,学生模型才能从教师模型的输出中学到更多有用的东西。当然, T T T 也不能过大,否则各类的预测值就都一样了。这里的温度 T T T 是一个需要手动调整的超参数。
蒸馏哪一层?
学生模型是否可能优于教师模型?
如果学生/教师模型规模/表达能力差距不太大,是有可能的。
一个原因是学生模型学习的是教师模型的输出,有更多有用的信息量。
可以看作是正则?
参数量化,通常在神经网络中,参数的数据类型是单精度的 fp32。如果我们对某些参数可以用更低的比特来表示,比如 fp16、int8,那么网络的规模明显会降低很多。
除了低比特表示之外,另一个做法是参数聚类(weight clustering)。对于模型中的各个参数,将他们数值接近的聚为一类,比如在下图中,不同的颜色就是不同一类的参数。然后,用一个 table 将这一类的参数数值的平均值记录下来。在网络中,记录每个参数属于哪一类,对于每一类,直接用 table 中存的平均值来表示。这样的话,在网络中存储每个参数所占用的空间就只与类的个数有关,比如图中有四个类,那么网络中每个参数就只需要 2 个位来存储。
在霍夫曼编码中,比较常出现的数据,用更少的位数来表示,比较少出现的数据,用更多的数据来表示,这样平均下来,会用到更少的存储空间。
参数量化的极致,当然就是将表示每个参数的所用的位数压到最低:一位。这样的话每个参数就只有两个可能值了。听起来很不可思议,这样的模型应该会效果很差甚至完全不 work 吧。但其实已经有一些研究表明,甚至有时二值网络的表现会比正常的网络表现还要好。一个解释是将网络二值化,相当于对网络参数进行了极致的正则化,要求每个参数只能是 -1/1,这样可以较好地缓解过拟合现象。
结构设计的模型压缩方法是指通过设计适合的网络模型结构,来得到更低参数量/计算量,效果接近的模型。
首先先来回顾一下标准卷积的参数量。一个标准卷积示意图如下,该层卷积的输入有 2 个通道,卷积核的 size 是 3x3,所以每个卷积核的参数量是 2x3x3,而一共有 4 个卷积核,即输出通道数为 4,因此总的参数量为:3x3x2x4=72。
然后来看一下深度可分离卷积是怎么做的。深度可分离卷积有两个步骤,第一个步骤是 Depthwise Convulution,即深度卷积,第二个步骤是 Pointwise Convolution,即点卷积。
Depthwise Convulution
在这一步中,有几个输入通道,就有几个卷积核,每个卷积核只处理对应的一个通道。比如在图中,输入有两个通道,那么卷积核的个数也有两个,分别处理一个输入通道,那么对应的,输出通道数也会是 2。也就是说,在 Depthwise Convlution 这一步 输入通道数=卷积核个数=输出通道数。
Depthwise Convlution 非常明显的缺陷是只能够处理同一张特征图内部的特征交互,而没有跨通道的特征交互,如果有某些 pattern 是跨通道才能识别到的,Depthwise Convlution 是无能为力的,这就需要 深度可分离卷积的第二个步骤:Pointwise Convlution。
Pointwise Convlution
经过 Depthwise Convlution 之后,特征的通道数与输入一致。在 Pointwise Convlution 中,是以 Depthwise Convlution 的输出作为输入。然后 Pointwise Convlution 是对 Deptwise Convlution 的输出再进行一次标准的卷积,但是有一点限制是卷积核的尺寸必须是 1x1。这样,Pointwise Convlution 只需要处理跨通道的信息交互就好了,不需要再考虑同一张特征图内部的特征交互,因为这在前一步中已经处理过了。
可以看到,深度可分离卷积两步完成之后输出特征图的尺寸是与标准卷积一致的,并且也都计算了特征图内部/跨通道之间的特征交互。因此,它的性能应该是不会差标准卷积太多的。现在,来计算一下这个例子中深度可分离卷积的计算量:Depthwise Convluton:3x3x2=18,Pointwise Convlution:1x1x2x4=8,共计 26。而刚才计算得到的标准卷积的计算量为 72。可以看到,深度可分离卷积对于模型参数量的压缩效果还是很可观的。
现在来计算以下一般情况下标准卷积与深度可分离卷积的参数量对比。
记输入输出通道数分别为 I , O I,\ O I, O,卷积核尺寸为 K K K ,
二者比值:
P
a
r
a
m
s
d
s
c
o
n
v
P
a
r
a
m
s
r
e
g
c
o
n
v
=
K
2
I
+
I
O
K
2
I
O
=
1
O
+
1
K
2
\frac{Params_{dsconv}}{Params_{regconv}}=\frac{K^2I+IO}{K^2IO}=\frac{1}{O}+\frac{1}{K^2}
ParamsregconvParamsdsconv=K2IOK2I+IO=O1+K21
在常见的深度 CNN 中,
O
O
O 是一个很大的大叔(128, 512, …),而卷积核尺寸通常为
K
=
3
K=3
K=3。因此,在常见的 CNN 中,换用深度可分离卷积,可粗略地估计卷积层的参数量较小到了原来的九分之一。
深度可分离卷积降低参数量的理论依据是低秩近似(Low Rank Approximation)。
先以线性层为例,如果有一个线性层,输入输出维度分别为 M , N M,\ N M, N ,那么它的参数量为: M × N M\times N M×N 。
如果想要降低它的参数量,可以怎么做呢?
做法就是在它们中间加一个映射到
K
K
K 维的中间层。现在一层变两层,参数量会变小吗?不难得到,变为两层之后,参数量为
M
×
K
+
K
×
N
M\times K+K\times N
M×K+K×N 。因此,如果
K
K
K 的值远远小于
M
,
N
M,N
M,N ,那么两层的参数量是会比一层的
M
×
N
M\times N
M×N 更小的。但是,做了这种变换之后,整个映射过程的秩是变小了,因此,该映射的表达能力可能有所降低。
实际上,深度可分离卷积的思想与上面线性层的低秩近似是类似的。也是通过将一层操作转换为两层操作,降低总的参数量,但是有最大限度地保证了整个映射的表达能力。如下图所示,在常规卷积和深度可分离卷积中,左上角的粉色元素的计算都是根据输入特征图中左上角的 18 个元素进行的。
与上面介绍的四种技术不同,动态计算的技术并不是直接压缩模型的大小,而是根据设备即时算力不同,动态地调整模型的大小。
比较常见的做法是在深度或宽度两个维度来动态调整模型的规模。
深度
我们知道,深度学习模型通常是有很多层堆叠起来的,然后用最后一层的输出特征接一个线性层,进行分类。在按深度动态计算中,我们可以在训练时给模型的每一层输出特征都接一个分类器,将他们的输出结果都与标签计算损失。这样就可以在推理时,根据当前的算力情况动态地选择经过多少层就输出预测结果。
这样做可能存在的问题是:一般CNN的浅层只需要提取一些局部的纹理特征,但是如果在浅层就要求模型能够根据特征进行预测,就强迫模型在浅层要提取一些高层的语义特征,这可能会影响最终的性能。
宽度
动态计算也可以在模型的宽度维度上来做,即根据当前算力选择用较宽或较窄的层。注意虽然模型的宽度可以选择,但是它们实际上是同一个模型,即图中相同颜色标记的权重都是一样的,只是动态地选择根据那些神经元进行计算。
之前介绍的动态计算方式都是需要根据即时算力来动态地选择模型的规模,实际上也可以根据样本的难度由模型自己决定在哪一层停止。比如图中的例子,比较简单的样本在浅层就已经有比较确定的答案,即可停止;而比较困难的样本在需要在深层才能做出预测。