不能。PyTorch prune模块不直接修改state_dict,而是通过hook和mask动态屏蔽权重;需调用prune.remove固化稀疏结构、转half、再保存state_dict才能减小体积。

PyTorch内置prune模块能直接改model.state_dict吗?
不能。调用 torch.nn.utils.prune.l1_unstructured 等函数后,模型参数不会被原地覆盖成稀疏张量;它只是在参数上叠加一个 forward_pre_hook,并在 module._parameters 中额外存一个 orig 和一个 mask。真正的权重仍是完整浮点张量,只是每次前向时按 mask 屏蔽部分值。
这意味着:剪枝后的 model.state_dict() 仍含全量参数,体积没变——你得手动固化稀疏结构。
- 用
torch.nn.utils.prune.remove(model, 'weight')永久移除 hook,把 mask 应用到原始权重上,再删掉orig和mask属性 - 这步必须在保存前做,否则
torch.save(model.state_dict(), ...)还是存了冗余字段 - 注意:remove 后无法再对同一参数调用 prune(会报错),如需迭代剪枝,应保留未 remove 的模型副本
剪枝后模型体积没变小?关键在保存前的三步操作
常见错误是直接 torch.save(model, path) 或只保存 state_dict() 却没 remove。实际压缩体积要靠「稀疏权重固化 + 权重转为半精度 + 去除冗余属性」。
- 先对所有已剪枝的层执行
torch.nn.utils.prune.remove(model, 'weight') - 再用
model.half()转float16(注意:确保推理环境支持 FP16,且 batch norm 层也要 .half()) - 最后保存时用
torch.save(model.state_dict(), path)—— 此时 dict 里只有被保留的权重,无 mask/orig 字段 - 验证体积:对比剪枝前后的
os.path.getsize(path),典型 ResNet18 在 50% 稀疏度下可减约 40%
global_unstructured 和 l1_unstructured 哪个更适合压缩体积?
global_unstructured 更适合——它跨所有指定层统一排序权重绝对值,全局裁掉最小的 N 个,避免某些层被过度剪、某些层几乎没动。而 l1_unstructured 是逐层独立剪,容易导致 head 层(如最后的 fc)被剪狠、backbone 层保留过多,最终稀疏分布不均,等效压缩率低。
立即学习“Python免费学习笔记(深入)”;
- 用法示例:
torch.nn.utils.prune.global_unstructured(..., pruning_method=torch.nn.utils.prune.L1Unstructured, amount=0.5) - amount=0.5 表示全局剪掉 50% 的参数,不是每层剪 50%
- 务必传入所有待剪枝参数的
(module, 'weight')元组列表,漏掉某层就等于放弃压缩那部分 - 对 Conv2d 和 Linear 层有效;BatchNorm2d 的 weight 不建议剪(影响归一化稳定性)
剪枝后精度暴跌?别跳过 fine-tuning 这一步
剪枝本质是破坏性操作,即使用 L1 排序,也会丢掉部分信息。零次微调(zero-shot pruning)通常会让准确率掉 5–15 个百分点,尤其在小数据集或高稀疏度(>60%)时更明显。
- 推荐方式:剪枝后用原始训练流程继续训练 5–10 个 epoch,学习率设为原训练的 1/10(例如 1e-4)
- 只更新未被剪掉的权重:PyTorch 默认已通过 mask 自动屏蔽梯度(
mask * grad),无需额外处理 - 如果用
remove固化了结构,后续训练就真只更新剩余权重;但若未 remove,hook 仍生效,训练中 mask 不变,仅更新 orig - 注意:BN 层统计量(running_mean / running_var)在剪枝后可能失准,fine-tuning 前建议用
model.train()并跑几个 batch 的 forward(不反传)来刷新
实际压缩效果取决于模型结构、任务难度和稀疏度设定。最容易被忽略的是:剪枝不是“一键瘦身”,而是「剪 → 微调 → 固化 → 保存」四步闭环;少任何一环,要么体积不变,要么精度崩盘。


















