Torch-Pruning

环境要求

• Python 环境（建议 3.8 及以上）
• PyTorch（README 测试覆盖 1.x 与 2.x，徽章标注 PyTorch>=2.0）
• NumPy（Torch-Pruning 仅依赖 PyTorch 和 Numpy）
• 可用的 pip 包管理工具

安装与启动步骤

1. 1安装 Torch-Pruning

   官方推荐方式，自动安装最新版；升级已有版本同样使用该命令。

   pip install torch-pruning --upgrade

   复制

2. 2源码可编辑安装

   需要改源码或跟最新代码时使用，克隆仓库后以 -e 模式安装到当前环境。

   git clone https://github.com/VainF/Torch-Pruning.git
   cd Torch-Pruning && pip install -e .

   复制

3. 3加载模型与示例输入

   以 torchvision 的 resnet18 为例准备模型和输入张量，作为剪枝的依赖图分析入口。

   import torch
   from torchvision.models import resnet18
   import torch_pruning as tp
   
   model = resnet18(pretrained=True)
   example_inputs = torch.randn(1, 3, 224, 224)

   复制

4. 4设置忽略层并建剪枝器

   分类器层不参与剪枝，需加入 ignored_layers；BasePruner 无需稀疏训练即可使用。

   ignored_layers = []
   for m in model.modules():
       if isinstance(m, torch.nn.Linear) and m.out_features == 1000:
           ignored_layers.append(m)
   
   pruner = tp.pruner.BasePruner(
       model,
       example_inputs,
       importance=imp,
       pruning_ratio=0.5,
       ignored_layers=ignored_layers,
       round_to=8,
   )

   复制

5. 5执行剪枝并统计

   调用 pruner.step() 完成剪枝，再用 count_ops_and_params 对比剪枝前后的 MACs 与参数量。

   base_macs, base_nparams = tp.utils.count_ops_and_params(model, example_inputs)
   tp.utils.print_tool.before_pruning(model)
   pruner.step()
   tp.utils.print_tool.after_pruning(model)
   macs, nparams = tp.utils.count_ops_and_params(model, example_inputs)
   print(f"MACs: {base_macs/1e9} G -> {macs/1e9} G, #Params: {base_nparams/1e6} M -> {nparams/1e6} M")

   复制

6. 6用自有代码微调

   README 未提供训练脚本，剪枝后需用你自己的训练代码微调剩余权重以恢复精度。

关键配置

如何确认成功

运行后打印的 MACs 与 #Params 相比剪枝前明显下降，且模型仍可正常前向推理。

常见问题

注意事项

• Torch-Pruning 仅依赖 PyTorch 和 Numpy，无需额外重型依赖。
• 通道维度建议对齐到 4x 或 8x，实际推理加速更明显。
• 剪枝后必须自行微调，README 未提供训练脚本。
• 交互式剪枝时 group 必须顺序处理，不要保存为列表后再统一执行。