PyTorch模型剪枝方法详解

本文深入解析了PyTorch模型剪枝的核心机制与实战陷阱，明确指出`prune.l1_unstructured`仅实施“软剪枝”——添加掩码置零但不减少参数量和计算开销，真正瘦身必须调用`prune.remove()`；针对边缘部署需求，详解如何用`prune.ln_structured`按通道进行结构化剪枝，并强调维度选择、范数计算及下游层适配的关键细节；同时预警保存加载时因未`remove`导致的`KeyError`、微调中因优化器状态错位引发的梯度爆炸等高频问题，提供可立即落地的规避策略与最佳实践，助你避开剪枝路上的绝大多数坑。

prune.l1_unstructured 为什么剪不掉参数量？

直接调用 prune.l1_unstructured 后发现模型 state_dict 大小没变、推理速度也没提升，是因为它只添加了「掩码（mask）」和「原始权重副本」，并未真正删除参数。PyTorch 的剪枝是“软剪枝”：被剪位置的权重在前向时被 mask 置零，但张量尺寸、内存占用、计算量全保持不变。

实操要点：

• 必须调用 prune.remove() 才能永久移除掩码、把零值对应参数真正丢弃（此时 weight 变成普通 nn.Parameter，不再含 _mask 和 _original_weight）
• prune.remove() 是不可逆操作，之后不能再对同一模块重复调用 prune.l1_unstructured
• 若想保留剪枝结构用于后续微调，就不要 remove；若要导出轻量模型或测推理延迟，必须 remove

如何对 Conv2d 层按通道（channel-wise）剪枝？

prune.l1_unstructured 是按元素（element-wise）剪，无法保证整个卷积核或通道被整块剔除——这对部署到边缘设备很不利，因为残存的零散非零值仍需访存和计算。真正常见的结构化剪枝需换方法：

• 用 prune.ln_structured，n=1 对应 L1-norm 按通道剪（dim=0 剪输出通道，dim=1 剪输入通道）
• 对 Conv2d，推荐 dim=0：剪掉整个输出通道（即删掉一个 out_channels 维度上的完整卷积核），下游层输入通道数会自动匹配
• 注意：ln_structured 要求传入的是「要剪的维度上可聚合的数值」，比如通道的 L1 norm 就是 torch.norm(weight, p=1, dim=[1,2,3])（对 [out_c, in_c, h, w] 的后三维度求 L1）

示例：

prune.ln_structured(module, name='weight', amount=0.3, n=1, dim=0)

剪枝后保存/加载模型报错 KeyError: 'xxx.weight_orig'

这是最常踩的坑：保存未 remove() 的剪枝模型，再加载时 PyTorch 试图还原带掩码的结构，但若加载代码里没重新注册剪枝器，就会找不到 _orig 或 _mask 属性。

• 方案一（推荐）：训练剪枝 → prune.remove() → 保存纯模型（torch.save(model.state_dict(), ...)）→ 加载时直接 load 到干净模型实例
• 方案二：保存整个模型对象（torch.save(model, ...)），但加载时必须确保环境完全一致（包括已 import prune、且模型类定义中已执行过相同 prune 操作），否则反序列化失败
• 切勿混合使用：比如保存前没 remove，加载后又手动调 prune.remove() —— 此时 weight_orig 已被删，会触发 AttributeError

剪枝 + 微调时 loss 突然爆炸或梯度为 nan

剪枝本身不改前向逻辑，但微调时梯度更新可能击穿被置零区域，尤其当用了 prune.random_unstructured 或未合理缩放学习率。

• 剪枝后首次微调，建议将学习率降到原训练的 1/5～1/10，避免权重在零附近剧烈震荡
• 确认优化器是否绑定了被剪枝的参数：如果对某层做了 prune.l1_unstructured 但没 remove()，优化器仍会更新 weight_orig，而前向走的是 masked weight，导致行为错位
• 更稳妥做法：剪枝 → prune.remove() → 用新 model.parameters() 重建 optimizer（丢弃旧状态），再微调

结构化剪枝后通道数变化，务必检查下游层（如 Linear 层）的 in_features 是否仍匹配——PyTorch 不会自动重算，这里容易静默出错。

今天关于《PyTorch模型剪枝方法详解》的内容介绍就到此结束，如果有什么疑问或者建议，可以在golang学习网公众号下多多回复交流；文中若有不正之处，也希望回复留言以告知！