📚 课程目录第4章剪枝工具链

4、剪枝工具链：PyTorch原生剪枝API、torch.nn.utils.prune详解、自定义剪枝函数

剪枝这件事，说白了就是给模型「瘦身」。但怎么剪、剪哪里、用什么工具剪，这里头门道不少。我刚开始做剪枝的时候，也踩过不少坑——比如用错了API，结果模型直接崩了。

这一章，咱们就聊聊PyTorch自带的剪枝工具链。你想想看，官方提供的API，稳定性和兼容性肯定是最好的。而且，学会了这些基础，后面自定义剪枝函数也就水到渠成了。

4.1 torch.nn.utils.prune：核心剪枝方法

PyTorch 在 torch.nn.utils.prune 里内置了好几种剪枝策略。我个人最常用的是下面这几个：

方法名	剪枝粒度	适用场景
`random_unstructured`	单个权重	快速验证、基线对比
`ln_structured`	通道/卷积核	结构化剪枝，硬件友好
`global_unstructured`	全局权重	按重要性全局排序剪枝
`custom_from_mask`	自定义	加载预计算好的 mask

关键点：这些方法都不会直接修改原始权重，而是创建一个同形状的 weight_mask 附加到模块上。前向传播时，权重 = 原始权重 × mask。这样做的好处是——你可以随时恢复被剪掉的连接。

举个例子，对卷积层做随机非结构化剪枝：

import torch.nn.utils.prune as prune

# 假设 model.conv1 是一个 Conv2d 层
prune.random_unstructured(
    module=model.conv1,
    name="weight",
    amount=0.3  # 剪掉 30% 的权重
)

# 查看剪枝后的效果
print(model.conv1.weight_mask)  # 0/1 组成的 mask 张量
print(model.conv1.weight)       # 注意：这里看到的是原始权重 × mask 的结果

嗯，这里要注意：amount 参数可以是浮点数（比例）或整数（个数）。我习惯用比例，因为不同层大小不一样，按比例更直观。

4.2 原生剪枝 API：检查、移除与恢复

剪完之后，怎么知道哪些层被剪过了？PyTorch 提供了几个实用工具：

is_pruned(module)：检查模块是否被剪枝过。
remove(module, name)：移除剪枝，恢复原始权重（但 mask 会被永久合并到权重中）。
compute_total_size(module)：计算剪枝后的参数总量。

我的习惯：每次剪枝后，我都会用 is_pruned 检查一遍，确保剪枝真的生效了。曾经有一次我忘了调用 apply，结果模型根本没变瘦，白忙活半天。

来看一个完整的流程：

# 1. 剪枝
prune.ln_structured(model.conv2, name="weight", amount=0.4, n=2, dim=0)

# 2. 检查
print(prune.is_pruned(model.conv2))  # True

# 3. 如果想永久化剪枝（不可逆）
prune.remove(model.conv2, "weight")
# 此时 model.conv2.weight 变成了剪枝后的权重，不再有 mask

# 4. 计算压缩率
original = model.conv2.original_weight.numel()
pruned = model.conv2.weight.numel()
print(f"压缩率: {1 - pruned/original:.2%}")

警告：remove() 操作是不可逆的！一旦调用，原始权重就被覆盖了。我建议在最终确定剪枝方案后再执行 remove，中间调试阶段保留 mask 更安全。

4.3 自定义剪枝函数：从入门到实战

内置方法虽然方便，但有时候满足不了我们的需求。比如，我想根据权重的梯度大小来剪枝，或者按某种自定义的评分标准。这时候就需要自己写剪枝函数了。

自定义剪枝的核心是继承 BasePruningMethod，实现两个方法：

compute_mask(t, default_mask)：根据张量 t 计算新的 mask。
apply(module, name)：将剪枝应用到模块上（通常直接复用父类方法）。

我给大家写一个实战例子——按权重值大小剪枝，但保留前 10% 的最大值：

import torch
import torch.nn.utils.prune as prune
from torch.nn.utils.prune import BasePruningMethod

class TopKPruning(BasePruningMethod):
    """保留权重值最大的 k 比例"""
    PRUNING_TYPE = 'unstructured'

    def __init__(self, keep_ratio):
        super().__init__()
        self.keep_ratio = keep_ratio

    def compute_mask(self, t, default_mask):
        # 计算值
        abs_t = torch.abs(t)
        # 找到阈值：保留比例对应的分位数
        k = int(t.numel() * self.keep_ratio)
        if k == 0:
            return torch.zeros_like(t, dtype=torch.bool)
        # 取前 k 大的索引
        _, indices = torch.topk(abs_t.view(-1), k)
        mask = torch.zeros_like(t, dtype=torch.bool)
        mask.view(-1)[indices] = True
        return mask

# 使用自定义剪枝
TopKPruning.apply(model.conv3, 'weight', keep_ratio=0.7)
# 保留值最大的 70% 权重，剪掉 30%

实战经验：我在 YOLOv5 的剪枝项目中，就用这种自定义方法做了「梯度感知剪枝」。具体来说，不是只看权重大小，而是结合了权重的梯度信息——梯度大的权重说明对 loss 影响大，应该优先保留。效果比纯按权重剪枝提升了 2-3 个 mAP 点。

4.4 剪枝后的模型导出与部署

剪枝完成后，别忘了做两件事：

调用 prune.remove()：把 mask 合并到权重中，否则导出的模型还带着 mask 结构。
用 torch.jit.script 或 torch.onnx.export 导出：确保剪枝后的稀疏结构能被推理框架正确加载。

我曾经踩过一个坑：剪枝后直接保存了 state_dict，结果加载时发现权重形状没变，但很多值变成了 0。后来才意识到，没有调用 remove 的话，保存的只是原始权重，mask 丢了。嗯，从那以后我养成了「剪枝 → 验证 → remove → 导出」的固定流程。

小技巧：对于非结构化剪枝（大量 0 值），可以考虑用稀疏张量格式存储，能进一步减少模型体积。PyTorch 的 torch.sparse 模块可以帮忙。

4.5 本章小结

这一章我们聊了 PyTorch 剪枝工具链的三个层次：

内置方法：random_unstructured、ln_structured、global_unstructured 等，开箱即用。
原生 API：is_pruned、remove、compute_total_size，帮你管理剪枝状态。
自定义函数：继承 BasePruningMethod，实现自己的剪枝逻辑。