A lightweight post-training pruning module built on top of PyTorch's Modules.
Part of the Tiny<X> series: TinyP (pruning), TinyQ
(quantization), TinyC (compiler), TinyI (interpreter).
TinyP implements the classic pruning pipeline — train → prune → fine-tune — picking up from a model that is already trained:
flowchart LR
A["Pretrained<br/>dense model"]
subgraph tinyp ["TinyP"]
direction LR
B["Prune<br/>Pruner.prune()"]
C["Fine-tune<br/>finetune()"]
B --> C
end
D["Compressed model<br/>sparse, or structurally smaller"]
A --> B
C --> D
C -.->|"iterate for higher sparsity"| B
- Prunes
nn.Conv2dandnn.Linearlayers of any pretrained PyTorch model - Pruning criterion — which weights to remove: weight magnitude
|W|, per-layer or global. First-order (Taylor) and second-order (Hessian) saliency are on the roadmap. - Pruning granularity — what structure the sparsity has: unstructured (fine-grained) and structured (filter / channel). Block-wise N:M is on the roadmap.
- Fine-tuning for accuracy recovery: masks are re-asserted after every optimizer step, so gradient updates cannot resurrect pruned weights
- Sensitivity scan: measure each layer's tolerance to sparsity, then assign per-layer ratios instead of a uniform one
- Model support: PyTorch models from Hugging Face Hub, torchvision, and any locally registered architecture
- Offline-first approach: no automatic downloads from the cloud
- Built-in benchmarking: sparsity, MACs, latency and memory footprint tracking
Every method is a choice on those two axes, and the granularity is what decides whether you get a speedup:
| Method | Criterion | Granularity | Tensor shapes | Speedup | Works on |
|---|---|---|---|---|---|
magnitude |
|W| | unstructured (fine-grained) | unchanged | none | any model |
structured_conv |
L2 norm per channel | structured (filter) | resized | real | plain conv chains |
structured_mask |
L2 norm per channel | structured, simulated | unchanged | none | any model |
lottery_ticket |
|W| + weight rewinding | unstructured | — | — | not yet implemented |
Unstructured pruning removes individual weights anywhere in the network. It reaches the highest compression for a given accuracy, but the resulting sparsity is irregular: dense CPU/GPU kernels cannot skip the zeros, and the cost of looking up sparse indices (indirection overhead) cancels what little is saved. The result is zero speedup without specialized sparse hardware.
Structured pruning removes whole filters, so the layer simply becomes a smaller dense layer. Standard hardware benefits immediately, with no indexing tricks — at the price of a lower achievable compression ratio. It only applies where conv layers form a straight chain; TinyP refuses residual, grouped and branching topologies rather than silently corrupting them.
That trade-off is the central one in pruning, and TinyP's benchmark reports both sides rather than picking the flattering number.
TinyP/
├── logs/ # Benchmark and training logs
├── models/ # Local model storage + reference architectures
│ └── vgg.py # VGG-CIFAR10 (MIT 6.5940 Lab 1)
├── scripts/
│ ├── fetch_assets.py # One-time asset download (only network I/O)
│ └── check_invariants.py # AST guards: offline-first, dropped security patterns
├── tests/ # pytest suite, incl. golden checks from the lab
├── BENCHMARK.md # How to run and read benchmarks
├── tinyp.py # Core pruning library
├── train.py # Sparsity-aware fine-tuning
├── utils.py # Metrics, model loading, inference helpers
├── examples.py # CLI
└── bench.py # Benchmarking tools
Note
TinyP is built with efficiency in mind to be used at the edge (locally) on both CPU and GPU based systems.
git clone https://github.com/diesimo-ai/TinyP.git
cd TinyP
# uv (recommended)
uv venv --python 3.12
uv pip install -e ".[profile,dev]"
# or pip
pip install -r requirements.txttorchprofile is optional — TinyP falls back to a built-in MAC counter that needs nothing
beyond torch.
Important
TinyP works in offline-mode only. Nothing reachable from import tinyp touches the
network; downloads are an explicit, separate command.
# MIT 6.5940 pretrained VGG on CIFAR-10, plus the dataset
python scripts/fetch_assets.py --allFor Hugging Face models, download them yourself:
huggingface-cli download --resume-download facebook/opt-125m --local-dir ./models/facebook/opt-125mfrom tinyp import Pruner
from utils import load_model
# Load model
model, _ = load_model("./models/vgg-cifar10/vgg.cifar.pretrained.pth", arch="vgg")
# Create tinyp pruner object
p = Pruner(model)
# Prune model (magnitude, structured_conv or structured_mask)
pruned = p.prune(sparsity=0.9, method="magnitude")
# Inspect what happened, layer by layer
for row in p.summary():
print(f"{row['name']:<28}{row['sparsity'] * 100:.2f}%")
# Save pruned model
p.export("./pruned_model.pth", pruned)Structured pruning needs one example input to trace the graph:
import torch
pruned = Pruner(model).prune(
sparsity=0.3,
method="structured_conv",
example_input=torch.randn(1, 3, 32, 32),
)python examples.py \
--model_path "./models/vgg-cifar10/vgg.cifar.pretrained.pth" \
--arch vgg \
--sparsity 0.9 \
--method magnitude \
--test_inference \
--pruned_path "./pruned_model.pth"| Argument | Default | Description |
|---|---|---|
--model_path |
None |
Local checkpoint or HF model directory |
--arch |
vgg |
Registered architecture name |
--sparsity |
0.9 |
Target sparsity — 0.9 prunes 90% |
--method |
magnitude |
magnitude, structured_conv, structured_mask, lottery_ticket |
--scope |
layer |
layer for a per-tensor threshold, global for one across the model |
--input_shape |
3,32,32 |
Example input shape, used to trace the graph |
--finetune_epochs |
0 |
Epochs of sparsity-aware fine-tuning (0 skips) |
--test_inference |
off | Run an inference smoke test after pruning |
--pruned_path |
./pruned_model.pth |
Where to write the pruned state dict |
python bench.py \
--model_path "./models/vgg-cifar10/vgg.cifar.pretrained.pth" \
--arch vgg --method structured_conv --sparsity 0.3 --evalBENCHMARK.md explains what each metric means, how to avoid the measurement mistakes that flatter a method, and what to record.
Aggressive pruning costs accuracy until the model is retrained. TinyP keeps the model sparse throughout — masks are re-asserted after every optimizer step:
from train import finetune
pruner = Pruner(model)
pruned = pruner.prune(sparsity=0.9, method="magnitude")
pruned, best_accuracy = finetune(pruned, train_loader, epochs=5,
pruner=pruner, test_loader=test_loader)Layers differ in how much pruning they tolerate. Scan first, then assign:
sparsities, curves = pruner.sensitivity_scan(eval_fn=lambda m: evaluate(m, test_loader))
pruned = pruner.prune(sparsity_dict={"backbone.conv7.weight": 0.9,
"backbone.conv0.weight": 0.0})See ROADMAP.md for the full plan — new pruning criteria, granularities and schedules, the benchmarking needed to compare them, and the deployment work that turns sparsity into an actual saving.
Nearest milestones:
torch.fxdependency graph, so ResNet / MobileNet can be structurally pruned instead of refused- Taylor (first-order) saliency as an alternative to magnitude
- N:M semi-structured sparsity (2:4)
- Iterative pruning, and
lottery_ticket - Sparse storage formats, so unstructured sparsity is a real file-size win
- Chaining into TinyQ for the full Deep Compression pipeline
Structured pruning changes the architecture itself. VGG's first conv, before and after
structured_conv at 30%:
(backbone): Sequential(
- (conv0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
- (bn0): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, ...)
+ (conv0): Conv2d(3, 45, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+ (bn0): BatchNorm2d(45, eps=1e-05, momentum=0.1, affine=True, ...)
(relu0): ReLU(inplace=True)The lab's original code does not show this — it warns that "Tensor's shape has been
changed, but the print info of model doesn't change", because weight.set_() leaves
out_channels stale. TinyP updates the module attributes too.
Fine-grained pruning leaves print(model) unchanged — the tensors keep their shape and
only their contents become sparse. Use pruner.summary() instead:
Layer Shape Sparsity
-------------------------------------------------------------
backbone.conv0.weight (64, 3, 3, 3) 89.99%
backbone.conv1.weight (128, 64, 3, 3) 90.00%
backbone.conv2.weight (256, 128, 3, 3) 90.00%
...
classifier.weight (10, 512) 90.00%
You can also use a tool like NETRON to compare both models in depth.
VGG-CIFAR10, the same pretrained model as the lab, measured on an RTX 4050 Laptop / i5-12450HX with torch 2.13. Dense baseline: 92.95% accuracy, 9.23 M params, 35.20 MiB, 606 M MACs, 5.9–6.3 ms CPU latency.
magnitude @ 0.9 |
structured_conv @ 0.3 |
|
|---|---|---|
| Sparsity | 89.95% | 0.00% (tensors resized instead) |
| Params | 9.23 M (1.00×) | 5.01 M (1.84×) |
| Size (nonzero) | 3.54 MiB (9.95×) | 19.13 MiB (1.84×) |
| MACs | 606 M (1.00×) | 305 M (1.98×) |
| CPU latency | 5.50 ms (1.06× — noise) | 3.84 ms (1.63×) |
| Accuracy, no fine-tune | 14.57% | 36.80% |
| Accuracy, 5-epoch fine-tune | 91.93% (−1.02 pts) | 92.35% (−0.60 pts) |
The headline: fine-grained pruning compresses ~10× but does not accelerate at all, while 30% channel pruning gives a real 1.63× speedup. Both recover to within ~1 point of dense accuracy after five epochs of fine-tuning.
Note
"Size (nonzero)" is the theoretical size if zeros were free — not what lands on disk. A
90%-pruned state_dict is still 35.24 MB dense; CSR gets it to 17.67 MB (2.0×) and a
bitmask plus packed values to 4.68 MB (7.53×). Sparse formats spend real bytes
recording where the nonzeros are, so the often-quoted 10× is an upper bound rather than
a delivered number. Structured pruning needs no special format at all: its checkpoint is
an ordinary dense file, 1.84× smaller. See docs/results.md.
How to run and interpret benchmarks: BENCHMARK.md. Full results, hardware details and methodology: docs/results.md.
Contributions are welcome! Please see the Contributing Guidelines.
CI runs lint, the full test suite on Python 3.10–3.12, and the architectural invariant
checks — CONTRIBUTING.md has the commands to reproduce it locally.
See CHANGELOG.md.
This project is licensed under the MIT License - see the LICENSE file for details.
This project started as a learning exercise from MIT 6.5940: TinyML and Efficient Deep Learning Computing by Prof. Song Han's Han Lab, which taught me the core concepts behind neural network pruning.
Special thanks to:
- Prof. Song Han and the MIT Han Lab team for making the course and labs openly available
- Lab 1 (Pruning), the reference implementation TinyP is refactored from — see my notes and lab solutions
- The lectures behind this work: L03 Pruning I and L04 Pruning II