Environment:VainF Torch Pruning PyTorch Python Core

Knowledge Sources	Torch-Pruning PyPI torch-pruning
Domains	Deep_Learning, Model_Compression
Last Updated	2026-02-08 12:00 GMT

Overview

Python >= 3.7 environment with PyTorch >= 2.0 and NumPy, providing the core runtime for the Torch-Pruning library.

Description

This environment defines the minimal requirements for running the torch-pruning package (v1.6.x). The library is a pure-Python PyTorch extension that performs structural pruning via dependency graph analysis. It requires PyTorch 2.0 or later for the core package, though the reproduction scripts support PyTorch as far back as 1.8.1. Python 3.7 through 3.11 are officially supported according to setup.py classifiers. NumPy is the only non-PyTorch dependency.

The package is CPU-compatible for all core operations (dependency graph building, pruning, importance estimation). CUDA is only required for benchmarking utilities (latency measurement, memory profiling).

Usage

Use this environment for all Torch-Pruning workflows: CNN pruning, Vision Transformer pruning, LLM structural pruning, object detection pruning, and sparse training. This is the mandatory base prerequisite; other environments layer additional dependencies on top of this one.

System Requirements

Category	Requirement	Notes
OS	Linux, macOS, Windows	OS-independent per setup.py classifiers
Python	>= 3.7	Classifiers list 3.7 through 3.11; CI tests on 3.9 and 3.10
Hardware	CPU (minimum)	CUDA GPU recommended for training and benchmarking

Dependencies

System Packages

No system-level packages are required for the core library.

Python Packages

torch >= 2.0
numpy (any version)

Optional Packages

timm (PyTorch Image Models) -- import-guarded in op_counter.py; enables FLOPs counting for timm-specific modules
torchvision -- required for example scripts and model loading, not for the core library

Credentials

No credentials or environment variables are required for the core library.

Quick Install

# Install core package from PyPI
pip install torch-pruning

# Or install with specific PyTorch version
pip install torch>=2.0 numpy torch-pruning

Code Evidence

Package dependency declaration from setup.py:9:

requirements = ["torch>=2.0", "numpy"]

Python version requirement from setup.py:42:

python_requires=">=3.7",

Optional timm import guard from torch_pruning/utils/op_counter.py:15-19:

try:
    import timm
    has_timm = True
except:
    has_timm = False

MultiheadAttention fallback for older PyTorch from torch_pruning/ops.py:228-231:

try:
    TORCH_MHA = nn.MultiheadAttention
except:
    TORCH_MHA = DummyMHA  # for pytorch w/o MultiHeadAttention

PyTorch version check for expand operation from torch_pruning/dependency/index_mapping.py:21:

if node.type == ops.OPTYPE.EXPAND and torch.__version__ >= "1.8":
    update_expand_index_mapping(graph, node)

Common Errors

Error Message	Cause	Solution
`ModuleNotFoundError: No module named 'torch_pruning'`	Package not installed	`pip install torch-pruning`
`AttributeError: module 'torch.nn' has no attribute 'MultiheadAttention'`	PyTorch version too old	Upgrade to PyTorch >= 1.9 (MHA was added in 1.9); the library falls back to a DummyMHA but MHA pruning will not work
`RuntimeError` during `DependencyGraph.build_dependency`	AutoGrad disabled (`torch.no_grad()` context active)	Remove `torch.no_grad()` when building the dependency graph; TP uses autograd for model tracing

Compatibility Notes

PyTorch 1.8.x: Supported by reproduction scripts (reproduce/requirements.txt specifies torch>=1.8.1), but the main package requires PyTorch >= 2.0.
PyTorch 2.6.0+: When loading pruned models, torch.load() requires weights_only=False because structural pruning changes the model architecture.
CI Coverage: Tests run against PyTorch 1.8.1, 1.11.0, 1.12.1, and 2.x (latest) on Python 3.9-3.10.
timm integration: The timm package is optional; FLOPs counting works without it but may not handle timm-specific custom modules.

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment