This CVPR tool actually deletes neurons instead of masking them

What it does

Torch-Pruning is a structural pruning framework for PyTorch that physically removes parameters rather than zeroing them out with masks. It supports a wide range of architectures, from classic CNNs to modern vision transformers, diffusion models, and large language models including DeepSeek-R1-Distill, Llama, and Qwen. The library offers both low-level pruning functions and high-level pruners for global, group-level, and interactive workflows.

The interesting bit

The core algorithm, DepGraph, treats the model as a computational graph and automatically discovers coupled layers—so if you prune an output channel from one convolution, it knows exactly which batch norms, activations, and subsequent layers must also be trimmed to keep the network valid. This automates the tedious bookkeeping that usually makes structural pruning error-prone.

Key highlights

• Backed by a CVPR 2023 paper and an ECCV 2024 follow-up (Isomorphic Pruning) for vision transformers and modern CNNs.
• Works off-the-shelf with Hugging Face, Timm, and Torchvision models, including YOLO, SAM, BERT, and Swin Transformers.
• Requires only PyTorch and NumPy; compatible with PyTorch 1.x and 2.x.
• Includes high-level pruners for global pruning, sparse training, and interactive pruning scenarios.

Caveats

• Building the dependency graph requires AutoGrad to be enabled, so model analysis cannot run inside torch.no_grad().
• Pruned models must be saved and loaded as whole model objects rather than state_dict snapshots, and PyTorch 2.6.0+ users need weights_only=False for loading.
• You must explicitly validate pruning groups before trimming; the library guards against over-pruning channels to zero, but only if you invoke the check.

Verdict

Worth exploring if you need to shrink production PyTorch models for latency or memory, especially when dealing with complex modern architectures where manual layer coupling would be painful. Skip it if you are only experimenting with unstructured sparsity or prefer simple magnitude-based masking.