Pruning
剪枝会从训练好的神经网络中移除重要性较低的权重、滤波器或层,以缩小模型并加快推理速度,然后通过微调恢复性能。
剪枝是一种模型优化技术,它会移除对已训练神经网络输出贡献很小的权重、神经元、滤波器、注意力头或整个层,从而得到更小、所需内存、计算资源和能耗更少的模型。在深度学习中,剪枝与模型剪枝是同义词。大多数网络都严重过度参数化:Han 等人通过训练、移除低幅值连接并重新训练其余部分,在 ImageNet上将 AlexNet 的参数量从 6,100 万降至 670 万,将 VGG-16 的参数量从 1.38 亿降至 1,030 万,同时准确率没有下降。在内存、功耗和推理延迟预算紧张的边缘 AI部署中,剪枝是常见步骤。
剪枝的工作原理#
每种剪枝方法都要决定移除哪些参数以及何时移除。重要性分数通常取权重幅值,因为接近零的权重对输出的影响最小;基于梯度的分数则能更直接地估计每个参数对损失函数的影响。一次移除很大比例的参数,通常比迭代剪枝损失更多准确率;迭代剪枝会先移除一小部分参数,再对剩余权重进行微调以恢复性能,并重复这一过程,直到达到模型大小或速度目标。彩票假设进一步指出,稠密网络中包含稀疏子网络;将这些子网络重置为原始初始化后单独训练,其准确率可与完整网络相当。
非结构化、结构化与半结构化剪枝#
非结构化剪枝会将张量中任意位置的单个权重置零。它能在给定准确率下达到最高稀疏度,但结果是一个不规则的稀疏矩阵;除非稀疏运行时会跳过零值,否则稠密的 GPU和 CPU内核仍会以完整代价处理它。结构化剪枝会从卷积神经网络或Transformer中移除整个滤波器、通道、注意力头或层,因此剪枝后的每一层都会变成更小的稠密层,能在任何硬件上更快运行,但每移除一个参数会损失更多准确率。PyTorch 的剪枝工具会将这两类剪枝应用为现有层上的掩码:
import copy
from torch import nn
from torch.nn.utils import prune
# Two copies of one convolution layer with 32 filters of shape 16x3x3
unstructured = nn.Conv2d(16, 32, kernel_size=3)
structured = copy.deepcopy(unstructured)
# Unstructured: zero the 50% of individual weights with the smallest magnitude
prune.l1_unstructured(unstructured, name="weight", amount=0.5)
# Structured: zero the 50% of whole filters (dim=0) with the smallest L2 norm
prune.ln_structured(structured, name="weight", amount=0.5, n=2, dim=0)
for name, layer in [("Unstructured", unstructured), ("Structured", structured)]:
zeros = (layer.weight == 0).float().mean()
empty_filters = (layer.weight.flatten(1).abs().sum(dim=1) == 0).sum()
print(f"{name}: {zeros:.0%} of weights zeroed, {empty_filters}/32 filters entirely empty")两层的权重都减少了 50%,但只有结构化剪枝后的层有 32 个滤波器中的 16 个完全为空;这些滤波器可以直接移除,留下一个含 16 个滤波器的稠密卷积层。
半结构化剪枝介于两者之间。NVIDIA 的 2:4 稀疏模式会将每四个连续权重中的两个置零,Ampere 及更新架构 GPU 上的稀疏 Tensor Core 会跳过这些权重,从而以一半时间完成相同的矩阵运算;NVIDIA 测得,在 A100 上使用 TensorRT 对稀疏 ResNeXt-101 进行端到端推理,性能最高提升 20%。在稠密硬件上,零散的零值本身不会带来速度提升:将 Ultralytics YOLOv5x 剪枝至 30% 全局非结构化稀疏度,使用其 prune() 工具且不进行微调后,COCO mAP50-95 从 0.507 降至 0.489,而 V100 上的推理时间仍为每张图像 5.2 毫秒。
实际应用#
当模型必须在受限硬件上运行时,剪枝最为重要。经过通道剪枝的目标检测网络能满足无人机、智能摄像头和手机的内存及功耗预算,无需连接云端即可实现实时推理,并将图像保留在设备上以保护数据隐私。在制造业检测中,同样的压缩方法让缺陷检测器能够在生产线旁的低成本边缘模块上运行。
大型语言模型也会进行剪枝。SparseGPT一次性将 OPT-175B 和 BLOOM-176B 剪枝至 60% 的非结构化稀疏度,无需重新训练,困惑度仅略有上升。NVIDIA 的 Minitron工作通过移除层或隐藏维度,将 Llama 3.1 8B 剪枝为 4B 模型,然后利用原始模型进行知识蒸馏以恢复性能。
剪枝与相关技术的比较#
剪枝通常与其他压缩方法结合使用;这个术语在神经网络之外还有更早的含义。










