Torch-Pruning 项目教程

2026-01-16 10:32:41作者：幸俭卉

项目介绍

Torch-Pruning（TP）是一个专门为结构化剪枝设计的库。与现有的框架（如 torch.nn.utils.prune）不同，TP 会物理地移除参数，并自动裁剪其他依赖层。TP 是一个纯 PyTorch 项目，支持 PyTorch 1.x 和 2.0 版本。它实现了内置的计算图追踪、依赖图（DependencyGraph）、剪枝器等功能，适用于各种结构化剪枝任务。

项目快速启动

安装

首先，通过以下命令安装 Torch-Pruning：

pip install torch-pruning

快速启动示例

以下是一个简单的示例，展示如何使用 Torch-Pruning 对 ResNet18 进行剪枝：

import torch
from torchvision.models import resnet18
import torch_pruning as tp

# 加载预训练的 ResNet18 模型
model = resnet18(pretrained=True)

# 创建一个示例输入
example_inputs = torch.randn(1, 3, 224, 224)

# 定义重要性标准
imp = tp.importance.GroupNormImportance(p=2)

# 初始化剪枝器
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.MetaPruner(
    model=model,
    example_inputs=example_inputs,
    importance=imp,
    pruning_ratio=0.5  # 移除50%的通道
)

# 执行剪枝
pruner.step()

# 保存剪枝后的模型
torch.save(model, 'pruned_resnet18.pth')

应用案例和最佳实践

结构化剪枝

结构化剪枝是一种移除模型中一组参数的技术，这些参数分布在不同的层中。由于层之间的依赖关系，这些参数必须同时移除以保持模型的结构完整性。Torch-Pruning 通过实现 DependencyGraph 来自动识别这些依赖关系，并收集剪枝组。

实际案例

以下是一个实际案例，展示如何对一个卷积层进行结构化剪枝：

import torch
from torchvision.models import resnet18
import torch_pruning as tp

# 加载预训练的 ResNet18 模型
model = resnet18(pretrained=True)

# 创建一个示例输入
example_inputs = torch.randn(1, 3, 224, 224)

# 构建依赖图
DG = tp.DependencyGraph()
DG.build_dependency(model, example_inputs=example_inputs)

# 获取剪枝组
group = DG.get_pruning_group(model.conv1, tp.prune_conv_out_channels, idxs=[2, 6, 9])

# 执行剪枝
if DG.check_pruning_group(group):
    group.prune()

# 保存剪枝后的模型
torch.save(model, 'pruned_resnet18_conv1.pth')