Torch-Pruning implementsDepGraph-based structural pruning for diverse models, providing high-level pruners and a dependency graph to prune correlated channels across layers. It supports examples from LLMs, SAM, ViT, and more; latest releases include code style enhancements and core dependency refactor.
Collecting history — the radar snapshots this repo daily. The trend line appears after 3 days of data (1 so far).
What it is
Torch-Pruning is a Python framework for structural pruning of neural networks using a graph-based algorithm called DepGraph. It aims to prune coupled parameters across layers rather than masking parameters, enabling pruning of a wide range of models including LLMs, vision transformers, and other architectures.
How it works
The repository centers on DepGraph to identify dependencies among layers during pruning. Users build a DependencyGraph from a model and an example input, then obtain pruning groups corresponding to layers or blocks. Pruning is performed by pruning groups, ensuring consistency across connected layers. It also provides utilities to scan groups across a model and to prune with specific indices. Example usage demonstrates constructing the graph, locating a pruning group, and applying prune().
Getting started
Install instructions are provided:
pip install torch-pruning --upgrade
For editable installation:
git clone https://github.com/VainF/Torch-Pruning.git
cd Torch-Pruning && pip install -e .
The README shows an end-to-end example of building the dependency graph and pruning a convolutional layer:
import torch
from torchvision.models import resnet18
import torch_pruning as tp
model = resnet18(pretrained=True).eval()
# 1. Build dependency graph for a resnet18. This requires a dummy input for forwarding
DG = tp.DependencyGraph().build_dependency(model, example_inputs=torch.randn(1,3,224,224))
group = DG.get_pruning_group( model.conv1, tp.prune_conv_out_channels, idxs=[2, 6, 9] )
# 3. Do the pruning
if DG.check_pruning_group(group):
group.prune()
It also shows how to save/load the pruned model and how to inspect pruning groups. Additionally, there are instructions for scanning all groups:
for group in DG.get_all_groups(ignored_layers=[model.conv1], root_module_types=[nn.Conv2d, nn.Linear]):
idxs = [2,4,6]
group.prune(idxs=idxs)
print(group)
Recent releases
Latest releases include:
- v1.6.1 (2025-09-07): enhance code style, standardize import organization, add comprehensive Google-style docstrings, enhance type hints.
- v1.6.0 (2025-07-03): Refactor the core Dependency module.
- v1.5.3 (2025-06-13): various fixes including BatchnormPruner behavior and related updates.
- Additional notes reference examples for LLMs and pruning enhancements through 2024–2025.
"Full Changelog" available via the project page changelog.
Traction
Stars: 3336 Forks: 385 Open issues: 340
Behind the repo
Not applicable: no linked startup or company page provided in the facts.
Caveats
License: MIT Creation date: 2019-12-15 Last push: 2025-09-07 The README notes compatibility with PyTorch 1.x and 2.x, and that Torch-Pruning relies on PyTorch and NumPy. No specific runtime caveats beyond standard dependencies are listed in the provided material.