Torch-Pruning 项目教程
项目介绍
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')
典型生态项目
DepGraph
DepGraph 是一个用于通用结构化剪枝的算法,它建模了结构化剪枝中的层依赖关系,实现了任意结构的剪枝。Torch-Pruning 是 DepGraph 的实现库。
相关论文
- 论文:DepGraph: Towards Any Structural Pruning
- 工程:https://github.com/VainF/Torch-Pruning
社区支持
Torch-Pruning 有一个活跃的社区,可以通过 GitHub Issues、Discord 或 WeChat 群组进行交流和获取帮助。
- Discord: 链接
- WeChat 群组: Group-2, Group-1 (500/500 FULL)
通过这些资源,用户可以获取最新的更新、教程和最佳实践。