pytorch模型剪枝后体积未减小,是因为剪枝仅置零权重并添加掩码(weight_orig/weight_mask),state_dict中张量尺寸不变;真正减小体积需调用prune.remove移除冗余参数,并配合结构化剪枝与稀疏感知部署后端(如onnx+tvm或tensorrt)。

PyTorch模型剪枝后为什么体积没变小?
剪枝操作本身只把权重置零,state_dict 里的张量尺寸完全不变,所以 torch.save() 保存的文件大小几乎不受影响。真正减小体积需要配合结构化稀疏导出或转成支持稀疏格式的部署后端(如 TorchScript + custom pass、ONNX + 稀疏算子、或 TensorRT 的稀疏插件)。
实操建议:
- 用
torch.nn.utils.prune.l1_unstructured或global_unstructured快速验证剪枝效果,但别直接拿它去测 APK 大小 - 若目标是移动端,优先走
torch.jit.trace+torch.jit.optimize_for_inference,再手动剔除已剪枝但未删除的参数(调用prune.remove) - 确认是否启用
prune.remove:它会把weight_orig和weight_mask删除,用剪枝后的weight替代原始参数——这步漏掉,模型仍带冗余 buffer
如何让剪枝后的模型在Android上真正加速?
PyTorch Mobile(libtorch_android)默认不利用稀疏性,0 值参与计算,速度反而可能下降。必须配合结构化剪枝 + 算子替换,或导出为 ONNX 后用 TVM / Core ML Tools / Qualcomm SNPE 等工具链做稀疏感知优化。
关键路径选择:
- 若坚持 PyTorch Mobile:只用通道级(channel-wise)结构化剪枝(如
ln_structured),然后用torch.jit.script导出,并在 Java 层调用前用model.eval()+torch.backends.quantized.engine = 'qnnpack'(对卷积层有隐式稀疏友好调度) - 更推荐路径:剪枝 →
torch.onnx.export→ 用onnx-simplifier清理图 → 输入到onnxruntime-mobile(开启ORT_ENABLE_EXTENDED_KERNELS编译选项才支持稀疏卷积) - 避免在
forward中动态判断 mask:移动端没有 runtime sparse kernel 支持,硬写if weight.abs().sum() == 0: skip反而增加分支开销
剪枝时 prune.global_unstructured 和 prune.ln_structured 怎么选?
前者按绝对值大小全局裁剪参数,适合快速压缩和精度探针;后者按通道/滤波器维度裁剪,保留模型结构完整性,是移动端落地的前提。
区别与取舍:
-
global_unstructured:剪枝率易控(比如直接设amount=0.5),但生成的是非结构化稀疏矩阵,ARM CPU 上无法跳过计算,GPU 更是完全无加速 -
ln_structured(常用n=1):按 L1 范数裁剪整个卷积核或全连接列,保证剪掉的是完整通道,后续可安全地移除对应输出通道及下一层输入通道(即“通道剪枝”),这是 NCNN / MNN 等推理引擎能识别并优化的粒度 - 混合策略可行:先用
global_unstructured找敏感层(看 accuracy drop),再对非敏感层用ln_structured做结构化裁剪
剪枝后精度掉太多,微调(fine-tuning)该怎么做?
剪枝本质是破坏原有权重分布,必须微调,但不能简单 resume 原始训练流程——已剪枝的参数需保持 mask 不变,否则剪枝失效。
正确做法:
- 微调前调用
prune.custom_from_mask或确保所有prune.*操作已应用且未remove,这样mask还在module上作为 buffer 存在 - loss backward 后,在
optimizer.step()前插入prune.custom_from_mask或手动将梯度乘以 mask(param.grad *= mask),防止梯度更新到被剪位置 - 学习率要比原始训练低 3–5 倍(例如从 1e-3 改成 2e-4),因为权重空间已被约束,大步长容易震荡
- 不要用
prune.remove直到微调结束——否则 mask 丢失,无法再做梯度掩码
prune.remove、结构化选择、部署后端适配、微调梯度掩码——四步缺一不可。Python免费学习笔记(深入):立即使用
在学习笔记中,你将探索 Python 的核心概念和高级技巧!











