PyTorch的torch.nn.utils.prune仅添加掩码（mask）并保留原始权重，不会真正移除参数，因此保存state_dict()时文件大小不变；要获得轻量化模型，需执行prune.remove()并重新导出结构化稀疏权重。
pytorch的`torch.nn.utils.prune`仅添加掩码（mask）并保留原始权重，不会真正移除参数，因此保存`state_dict()`时文件大小不变；要获得轻量化模型，需执行`prune.remove()`并重新导出结构化稀疏权重。
在PyTorch中实现模型剪枝（如使用prune.global_unstructured）时，一个常见误区是：剪枝后直接保存model.state_dict()，却发现文件体积与原始模型几乎完全一致。这并非Bug，而是PyTorch剪枝机制的设计使然——它采用非侵入式、可逆的掩码（masking）策略，而非物理删除权重。
? 剪枝的本质：掩码 + 原始参数共存
调用prune.global_unstructured(...)后，PyTorch会对目标参数（如conv.weight）进行以下原地修改（in-place）：
- 新增一个缓冲区（buffer）：weight_mask（二值掩码，形状同weight）；
- 将原参数重命名为weight_orig（仍为Parameter类型，完整保留）；
- 创建一个计算图感知的张量 weight：其值 = weight_orig * weight_mask，用于前向传播。
这意味着：
- 模型参数总量未减少（weight_orig + weight_mask 甚至略微增大）；
- model.state_dict() 默认只保存 weight_orig 和 weight_mask，不包含已应用掩码的“有效权重”；
- prune.remove() 是关键一步：它将 weight_orig * weight_mask 固化为新的weight参数，并删除_orig和_mask，实现真正的结构化稀疏。
✅ 正确做法：剪枝 → 固化 → 保存
您原代码中虽调用了 prune.remove(...)，但存在两个关键问题：
- modules_list 构造错误：filter + map 后得到的是 (module, 'weight') 元组，而 prune.remove() 要求传入被剪枝的模块对象本身（即 module），而非元组；
- prune.remove() 执行时机与范围：应在所有剪枝完成后，对每个被剪枝的模块单独调用，且必须确保该模块确实已被剪枝（即存在 _orig 属性）。
以下是修正后的完整流程（含ResNet18示例）：
import torch
import torch.nn as nn
from torch.nn import utils as prune
def random_unstructured_pruning(pruning_rate: float, device, log_file):
# 加载原始模型
trained_model_path = "./models/trained_model.pth"
model = ResNet18(num_classes=10, input_channels=1).to(device)
model.load_state_dict(torch.load(trained_model_path, map_location=device))
# 收集所有支持剪枝的层（Conv2d/Linear/BatchNorm2d 的 weight）
modules_to_prune = []
for name, module in model.named_modules():
if isinstance(module, (nn.Conv2d, nn.Linear, nn.BatchNorm2d)):
if hasattr(module, 'weight') and module.weight is not None:
modules_to_prune.append((module, 'weight'))
# 全局非结构化剪枝（L1准则，剪除80%）
prune.global_unstructured(
modules_to_prune,
pruning_method=prune.L1Unstructured,
amount=pruning_rate # e.g., 0.8
)
# ✅ 关键修正：对每个被剪枝模块调用 prune.remove()
# 注意：prune.remove() 接收 module 和 param_name（字符串），不是元组
for module, param_name in modules_to_prune:
if hasattr(module, param_name + '_orig'): # 确保已剪枝
prune.remove(module, param_name)
# ✅ 此时 model.state_dict() 已不含 _orig/_mask，weight 为稀疏后固化值
pruning_rate_str = f"{int(pruning_rate * 100):02d}"
save_path = f"{model_saving_path}pruned_{pruning_rate_str}.pth"
torch.save(model.state_dict(), save_path)
# ✅ 验证：打印剪枝后各层非零权重比例
total_params = 0
nonzero_params = 0
for name, param in model.named_parameters():
if 'weight' in name:
total_params += param.numel()
nonzero_params += torch.count_nonzero(param).item()
sparsity = 1 - nonzero_params / total_params
print(f"Actual sparsity: {sparsity:.3f} ({nonzero_params}/{total_params} non-zero)")
⚠️ 重要注意事项
- prune.remove() 不可逆：执行后无法再通过prune.reverse()恢复剪枝状态，请在调试阶段谨慎使用；
- 保存完整模型 vs state_dict：若需部署到其他框架（如ONNX），建议先remove再导出；若仅用于PyTorch推理，也可直接保存整个模型（torch.save(model, ...)），但体积仍含冗余（因_orig/_mask可能残留）；
- BatchNorm层剪枝需谨慎：剪枝weight对BN影响较小，但剪枝bias或running_mean/var可能导致精度崩塌，通常仅剪枝Conv2d/Linear的weight；
- 量化与剪枝协同：真实压缩需结合量化（如torch.quantization）或稀疏格式存储（如CSR），单纯剪枝+state_dict保存仅减少计算量，不显著减小磁盘体积（浮点数仍占4字节）。
? 总结
PyTorch剪枝不是“删权重”，而是“加开关”。要获得真正轻量化的模型文件，必须：
- 使用 prune.global_unstructured 或 prune.ln_structured 等方法施加掩码；
- 显式调用 prune.remove(module, 'weight') 固化稀疏权重；
- 保存 model.state_dict() —— 此时weight已是稀疏张量，_orig/_mask已被清除。
只有完成这三步，.pth 文件大小才会随剪枝率显著下降，模型推理延迟与内存占用才能切实优化。
