4. 模型剪枝与稀疏化:结构化剪枝、非结构化剪枝、稀疏矩阵运算加速
模型剪枝,说白了就是给神经网络“瘦身”。
你想想看,一个训练好的模型里,很多参数其实都是冗余的。它们对最终结果贡献极小,甚至没有贡献。把这些参数去掉,模型变小了,推理速度自然就上来了。
我在项目中遇到过最极端的情况:一个语音唤醒模型,剪掉 80% 的参数后,精度只掉了 0.3%。但推理速度提升了将近 4 倍。嗯,这就是剪枝的魅力。
4.1 剪枝的核心思想
剪枝的本质,是寻找并移除那些“不重要”的权重。那怎么判断一个权重重不重要?
最常见的方法就是看权重的绝对值大小。绝对值越小,说明它对最终输出的影响越小,剪掉它影响也最小。
当然,也有更高级的方法,比如基于梯度、基于损失函数变化等。但说实话,在实际工程中,基于幅度的剪枝(Magnitude-based Pruning)是最简单、最稳定的。我个人的习惯是,先从这个方法入手。
4.2 非结构化剪枝:精细但“难用”
非结构化剪枝,就是对每个权重单独判断。绝对值低于阈值的,直接置为 0。
这样做的好处是精度损失小。因为你可以非常精细地控制哪些权重被剪掉。理论上,它能达到最高的压缩率。
但问题也来了。剪完之后,模型里的权重变成了稀疏矩阵。里面全是 0 和少数非 0 值。这种不规则的结构,对硬件非常不友好。
我曾经在一个 ARM 平台上试过。模型大小确实变小了,但推理速度反而变慢了。为什么?因为 CPU 没法高效地处理这种随机分布的稀疏数据。它需要大量的判断和跳转指令,反而拖慢了速度。
4.3 结构化剪枝:工程首选
结构化剪枝就不一样了。它不剪单个权重,而是剪整个“结构”。比如一个卷积核(Filter)、一个通道(Channel)、甚至是一整层(Layer)。
这样做的好处是,剪完之后模型结构依然是规则的。你可以直接得到一个更小的、稠密的模型。不需要任何特殊的硬件或库支持,就能获得实实在在的加速。
在语音算法里,我用的最多的是通道剪枝。比如对卷积层,我们计算每个输出通道的重要性(比如通过 BN 层的缩放因子 gamma),然后把不重要的通道整个去掉。
| 对比维度 | 非结构化剪枝 | 结构化剪枝 |
|---|---|---|
| 剪枝粒度 | 单个权重 | 通道/卷积核/层 |
| 精度损失 | 较低 | 相对较高 |
| 硬件加速 | 需要专用库 | 直接生效 |
| 工程落地 | 困难 | 容易 |
| 适用场景 | 学术研究、极致压缩 | 工业部署、工程优化 |
你看这个表格就清楚了。对于咱们做语音算法落地的,结构化剪枝是更务实的选择。
4.4 稀疏矩阵运算加速
如果你非要走非结构化剪枝这条路,那你就得面对稀疏矩阵。
稀疏矩阵的加速,核心就是“只算非零元素”。怎么做到?主要有两种格式:
- CSR (Compressed Sparse Row): 按行压缩。用三个数组分别存储非零值、列索引、行偏移。
- CSC (Compressed Sparse Column): 按列压缩。原理类似,只是按列来。
下面是一个简单的 CSR 格式示例。假设我们有一个 4x4 的稀疏矩阵:
原始矩阵:
[[0, 0, 3, 0],
[1, 0, 0, 0],
[0, 0, 0, 5],
[0, 2, 0, 0]]
CSR 表示:
values = [3, 1, 5, 2] # 非零元素
col_indices = [2, 0, 3, 1] # 每个非零元素所在的列
row_ptr = [0, 1, 2, 3, 4] # 每行第一个非零元素在 values 中的起始索引
用 CSR 格式做矩阵乘法时,我们只遍历 values 数组,跳过所有零。这样计算量就从 O(n²) 降到了 O(nnz),其中 nnz 是非零元素个数。
4.5 剪枝流程与策略
剪枝不是一步到位的。我通常采用“训练-剪枝-微调”的迭代流程:
- 预训练: 先训练一个性能良好的大模型。
- 剪枝: 按照一定比例(比如 20%)剪掉不重要的权重或通道。
- 微调: 用原始训练数据,对剪枝后的模型进行少量训练,恢复精度。
- 重复: 回到步骤 2,逐步增加剪枝比例,直到精度出现不可接受的下降。
这里有个关键点:剪枝比例怎么定?我个人的经验是,不要一次性剪太多。比如目标压缩率是 50%,我会分 5 次,每次剪 10%,然后微调。这样精度损失会小很多。
下面这张图展示了整个剪枝与稀疏化的知识体系:
4.6 实战中的避坑指南
最后,分享几个我踩过的坑:
- 不要只看模型大小: 模型变小了不代表推理变快了。一定要上板子实测。我曾经有个模型剪到原来的 1/4,但推理时间只少了 10%。因为瓶颈在内存带宽,不在计算量。
- 微调学习率要调小: 剪枝后的模型已经很脆弱了。用原始学习率去微调,很容易把精度搞崩。我一般会把学习率降到原来的 1/10 甚至 1/100。
- 注意剪枝后的层间匹配: 结构化剪枝时,如果你剪掉了某个卷积层的通道,那下一层的输入通道数也要跟着变。这个逻辑一定要在代码里处理好,否则模型结构就乱了。
嗯,关于剪枝和稀疏化,今天就聊这么多。记住,没有最好的方法,只有最适合你场景的方法。动手试试看吧。