PyTorch剪枝后变慢是因为未调用prune.remove,仅加掩码未删除参数;TensorFlow需训练时剪枝+strip_pruning再转TFLite;ONNX无原生剪枝,须回退框架处理。

PyTorch里用torch.nn.utils.prune剪枝后模型变慢?先确认是否真的“剪掉了”
剪枝后推理没变快,大概率是只调用了prune.l1_unstructured这类函数,但没调用prune.remove。PyTorch的剪枝默认是“掩码式”的:参数还在,只是乘了个0掩码,前向时仍要加载、参与计算,甚至可能因掩码分支拖慢速度。
实操上必须分两步:
-
prune.l1_unstructured(model.layer, name='weight', amount=0.2)—— 加掩码 -
prune.remove(model.layer, 'weight')—— 永久删掉被剪参数,替换为nn.Parameter张量
漏掉第二步,model.state_dict()里还能看到weight_orig和weight_mask,说明剪枝没落地。此时model.layer.weight仍是带掩码的PruningContainer,不是普通Tensor。
TensorFlow/Keras想用tfmot.sparsity但导出TFLite报错Op type not registered 'SparseToDense'
这是典型训练后剪枝(post-training pruning)未适配TFLite流程导致的。TF-MOT的剪枝器在训练中插入了稀疏操作,但TFLite不支持运行时稀疏张量,导出时会卡在SparseToDense这类算子上。
立即学习“Python免费学习笔记(深入)”;
正确路径是走“训练时剪枝(pruning during training)+ 量化感知训练(QAT)+ strip_pruning”三步:
- 用
pruning_schedule(如PolynomialDecay)在训练中逐步施加剪枝 - 训练完调用
tfmot.sparsity.keras.strip_pruning(model)—— 它会移除所有剪枝相关层,只保留稀疏后的权重 - 再用
tf.lite.TFLiteConverter.from_keras_model转换,就不会遇到稀疏算子了
注意:strip_pruning只能在训练完成、评估达标后调用,不能在训练中途用,否则会破坏剪枝调度逻辑。
剪枝后精度掉太多,不是因为剪多了,而是没重训练或没调学习率
直接剪掉20%权重再跑几轮fine-tune,精度崩盘很常见。原因不是剪枝本身激进,而是剪枝改变了梯度流分布,原学习率对剩余参数来说往往太大。
关键调整点:
- fine-tune阶段学习率建议设为原训练的
0.01~0.1倍(比如原用1e-3,现在用1e-4) - 至少跑
5~10个epoch,且用带早停的验证集监控,别只看训练loss - 如果用结构化剪枝(如剪整个卷积通道),务必检查后续层的输入channel数是否自动对齐——有些框架不会自动改,得手动同步
in_channels
另外,l1_unstructured适合微调后剪,ln_structured(n≥2)更适合从头训练时嵌入,二者对精度影响模式不同,混用容易失配。
ONNX模型怎么剪枝?没有现成API,得回退到PyTorch或TF处理
ONNX是中间表示,本身不支持剪枝操作。所有所谓“ONNX剪枝工具”(比如onnxruntime.transformers里的某些函数)其实只是做等效替换或图优化,并不真正删参数。
真实可行路径只有两条:
- 在PyTorch里剪完、
strip_pruning、保存为.pt,再用torch.onnx.export导出新ONNX - 或用
onnx-simplifier清理冗余节点(比如Constant+Mul变成零张量),但这不算剪枝,只是图压缩
想绕过框架直接改ONNX权重?可以,但要手动遍历model.graph.initializer,把对应tensor.data改小、再更新tensor.dims,还要同步改所有引用它的NodeProto.input。出错概率高,不推荐。


















