跳转到主内容
websoft网络软件专家 - 深耕网络技术,打造实用软件!

PyTorch模型权重如何压缩_通过剪枝技术去除接近零的权重数值

剪枝前需先分析权重分布并谨慎选择剪枝层,剪枝后必须固化掩码、导出稀疏格式或ONNX并启用稀疏优化,且务必微调以恢复精度,否则部署时剪枝无效。 剪枝前先确认哪些权重真该删 PyTorch 里没有“自动识别冗余权重”的魔法函数,
torch.nn.utils.prune.l1_unstructured
这类方法只按数值大小排序,不看梯度、结构或实际贡献。直接对
weight
张量做全局阈值剪枝,常导致精度断崖式下跌——因为很多接近零的权重在训练中起着稳定梯度或补偿量化误差的作用。 实操建议: 先用
torch.histc(weight.abs(), bins=50)
看分布,别盲目设
amount=0.3
;若尾部密度极低(比如最后 10% 的值集中在 1e-6 量级),才考虑剪 优先剪
nn.Linear
和
nn.Conv2d
的
weight
,避开
BatchNorm2d
的
running_mean
或
bias
(剪了会破坏归一化) 剪完立刻用
model.apply(torch.nn.utils.prune.remove)
把掩码固化成 0,否则推理时仍参与计算 剪枝后模型体积没变小?那是没导出稀疏格式 PyTorch 默认保存的是稠密张量:即使你用
prune.custom_from_mask
把 90% 权重设为 0,
torch.save(model.state_dict(), 'x.pth')
仍存满 32-bit 浮点数,文件大小几乎不变。 实操建议: 导出前调用
torch.nn.utils.prune.remove(model, 'weight')
,再用
torch.save(..., _use_new_zipfile_serialization=False)
避免 ZIP 压缩干扰体积判断 真正减体积得转稀疏格式:
weight.to_sparse()
后保存,但注意:CUDA 不支持稀疏卷积推理,CPU 上也要用
torch.sparse.mm
替代常规
@
更实用的路径是导出为 ONNX,用
onnxruntime
的
OrtSession
加载时启用
sparse_optimization=True
,它会自动跳过零值计算 剪枝后精度掉太多,大概率是没重训练 一次性剪掉 40% 权重却不微调,相当于把模型硬掰弯后不校准——
prune.global_unstructured
只改数值,不更新其余参数的补偿能力。常见现象是验证集 top-1 准确率跌 5%+,且 loss 曲线震荡剧烈。 实操建议: 剪枝后至少跑 3–5 个 epoch 的 fine-tuning,学习率设为原训练的 1/10(如
1e-4
),优化器用
SGD
比
Adam
更稳 别在剪枝后的模型上继续剪:第二次
prune.l1_unstructured
会作用于已含大量零的张量,容易误删关键非零权重 用
torch.nn.utils.prune.is_pruned(layer)
检查是否重复应用,避免掩码叠加 TensorRT 或 TFLite 部署时剪枝失效 这些推理引擎加载 PyTorch 模型时,通常走 ONNX 中间层,而 ONNX 导出默认忽略
prune
的掩码逻辑——
weight
张量被复制为稠密形式,零值照常参与计算,白剪。 实操建议: 导出 ONNX 前务必执行
torch.nn.utils.prune.remove(model, 'weight')
,确保
state_dict
里
weight
已是纯 0/非0 稠密张量 ONNX 导出时加参数
do_constant_folding=True
,让零值乘法被折叠掉(但仅限静态图) TensorRT 8.6+ 支持稀疏 GEMM,但需手动指定
builder_config.set_flag(BuilderFlag.SPARSE_WEIGHTS)
,且只对
Conv
和
FC
层生效 真正难的不是剪多少,而是剪完哪一层、在哪一步固化、在哪个后端能真正跳过计算——这些细节错一个,前面全白忙。

相关文章