第3章模型剪枝实战

第3章：模型剪枝实战——基于PyTorch的剪枝工具使用

模型剪枝，说白了就是把神经网络里那些“不干活”的参数给干掉。我刚开始接触剪枝时，总觉得这玩意儿很玄乎——把参数删了，模型还能正常工作？后来亲手试了一次，才发现这其实是个“去芜存菁”的过程。

PyTorch从1.4版本开始，内置了torch.nn.utils.prune模块。我个人习惯用这个工具做快速原型验证，它最大的好处是：不需要改模型结构，直接在原模型上做“手术”。

3.1 剪枝的核心思想

剪枝的本质，是找到那些对输出贡献很小的权重，把它们置零。你想想看，一个卷积核有9个参数，其中3个值接近0，那这3个参数基本就是“摆设”。

剪枝分为两种：

非结构化剪枝：单个权重置零，稀疏但零散
结构化剪枝：整通道/整层删除，规整但粗暴

我在项目中遇到过这样的情况：非结构化剪枝虽然压缩率高，但硬件加速器不支持稀疏矩阵运算，结果推理速度反而没提升。所以选哪种剪枝方式，得看你的部署平台。

核心原则：剪枝不是目的，加速才是。别为了剪枝而剪枝。

3.2 PyTorch剪枝工具概览

torch.nn.utils.prune提供了三类核心功能：

功能类别	函数/类	说明
剪枝方法	prune.L1Unstructured	按L1范数剪枝（最常用）
剪枝方法	prune.RandomUnstructured	随机剪枝（用于对比实验）
剪枝方法	prune.ln_structured	结构化剪枝（按通道）
应用/移除	prune.remove	永久化剪枝掩码
全局剪枝	prune.global_unstructured	跨层统一剪枝比例

嗯，这里要注意：PyTorch的剪枝是“掩码式”的——它不会真的删除参数，而是生成一个和二值掩码，把对应权重置零。这样做的好处是方便恢复，坏处是模型文件大小不变。

3.3 实战：对卷积层做L1非结构化剪枝

咱们直接上代码。假设你有一个训练好的LeNet模型，要对第一个卷积层剪掉20%的权重。

import torch
import torch.nn.utils.prune as prune
import torch.nn.functional as F

class LeNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(1, 6, 5)
        self.conv2 = torch.nn.Conv2d(6, 16, 5)
        self.fc1 = torch.nn.Linear(256, 120)
        self.fc2 = torch.nn.Linear(120, 84)
        self.fc3 = torch.nn.Linear(84, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.fc3(x)

model = LeNet()
# 加载预训练权重... 这里省略

# 对conv1的权重做L1非结构化剪枝，剪掉20%
prune.l1_unstructured(model.conv1, name='weight', amount=0.2)

# 查看剪枝后的权重
print(model.conv1.weight)
# 你会发现很多权重变成了0

这段代码跑完后，model.conv1.weight里大约20%的参数变成了0。但注意，model.conv1现在多了一个属性：weight_mask，里面存的就是掩码。

我的小技巧：剪枝后一定要做微调（fine-tune）。我曾经偷懒没微调，结果模型精度掉了5个点。微调3-5个epoch就能恢复大部分精度。

3.4 结构化剪枝：按通道剪

非结构化剪枝生成的稀疏矩阵，在通用硬件上很难加速。结构化剪枝就不一样了——直接干掉整个通道，模型体积和计算量都实打实地减少。

# 对conv1按通道剪枝，剪掉2个输出通道
prune.ln_structured(model.conv1, name='weight', amount=2, n=2, dim=0)

# dim=0 表示输出通道维度
# n=2 表示使用L2范数评估通道重要性

这里有个坑：结构化剪枝后，下一层的输入通道数也要对应调整。PyTorch的prune模块不会自动帮你做这个调整，你得手动改网络结构。所以，结构化剪枝更适合用在你自己设计的、可灵活调整宽度的网络上。

避坑指南：我曾经在ResNet上做结构化剪枝，剪完才发现shortcut连接的通道数对不上了。后来我改用torch-pruning这个第三方库，它自动处理通道对齐。如果你用PyTorch原生工具，一定要自己检查维度。

3.5 全局剪枝：跨层统一比例

有时候你想让整个模型达到某个稀疏度，而不是每层单独剪。这时候用prune.global_unstructured。

# 收集所有要剪枝的参数
parameters_to_prune = [
    (model.conv1, 'weight'),
    (model.conv2, 'weight'),
    (model.fc1, 'weight'),
    (model.fc2, 'weight'),
    (model.fc3, 'weight'),
]

# 全局剪枝，整体稀疏度50%
prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.5,
)

全局剪枝的好处是：不重要的层多剪点，重要的层少剪点。我做过实验，同样50%稀疏度，全局剪枝比逐层剪枝精度高1-2%。

3.6 剪枝后的模型处理

剪枝完成后，模型里还留着掩码。如果你想导出模型做推理，需要做两件事：

移除掩码：用prune.remove()把掩码合并到权重中
重新打包：用torch.jit.script或torch.onnx.export导出

# 永久化剪枝
prune.remove(model.conv1, 'weight')
# 此时model.conv1.weight就是置零后的权重，不再有weight_mask属性

# 导出为TorchScript
scripted_model = torch.jit.script(model)
scripted_model.save('pruned_lenet.pt')

嗯，这里要注意：prune.remove()是不可逆的。一旦移除，你就没法恢复被剪掉的权重了。所以建议先保存一份剪枝前的模型。

3.7 剪枝效果评估

剪枝完，你得看看效果。我一般从三个维度评估：

指标	计算方法	我的经验阈值
稀疏度	零值参数 / 总参数	50%-80% 比较安全
精度损失	剪枝前精度 - 剪枝后精度	≤1% 可接受
推理加速比	剪枝前耗时 / 剪枝后耗时	≥1.5x 才算有效

我见过有人剪到90%稀疏度，精度只掉0.5%，但推理速度没变——因为硬件不支持稀疏加速。所以，一定要在目标硬件上测实际速度。

3.8 本章知识体系

下面这张图总结了剪枝的核心流程和工具链：

这张图把剪枝流程串起来了：从原始模型出发，选方法、应用剪枝、微调、评估、导出。每一步都有对应的PyTorch函数。

3.9 常见问题与避坑

最后，分享几个我踩过的坑：

剪枝后精度暴跌：别一次性剪太多。我建议从10%开始，逐步增加，每次剪完都微调一下。
模型文件没变小：PyTorch的剪枝只是置零，没删除参数。想真正减小文件，得用torch.save时只保存非零参数，或者转成稀疏格式。
推理速度没提升：检查你的硬件是否支持稀疏矩阵运算。如果不支持，考虑结构化剪枝或改用torch.sparse。
梯度传播异常：剪枝后的权重是0，但梯度依然会回传。如果你不想让被剪掉的权重再更新，可以在优化器里设置params时排除它们。

嗯，剪枝这块内容就这些。工具本身不复杂，复杂的是怎么剪得恰到好处。多试几次，你就能找到感觉。

📚 课程目录