latest (dev)
Copy
Latest development documentation · Updated 2026-10-08
Pruning - tensorplay.ao.pruning
tensorplay.ao.pruning sparsifies a trained model: it zeroes weights that
contribute least, so the model stays the same size but computes less (with
sparse-aware kernels) or compresses better on disk. Pruning is implemented
as a reparameterization — the pruned weights are replaced by a mask
applied by a forward pre-hook, so the sparsity pattern stays inspectable
and reversible until you explicitly make it permanent.
Two families of methods are provided:
Unstructured — zeroes individual weights chosen by importance:
l1_unstructured(smallest magnitude),random_unstructured(uniformly random),custom_from_mask(your own mask).Structured — removes whole channels along a dimension, producing layout-friendly sparsity:
ln_structured(smallest n-norm channels),random_structured.
Both take (module, name, amount, ...): the module, the name of the
parameter or buffer to prune (e.g. "weight"), and how much to prune —
an int (number of entries/channels) or a float (fraction of them).
import tensorplay as tp
from tensorplay.ao.pruning import l1_unstructured, ln_structured, is_pruned
model = tp.nn.Linear(64, 32, bias=False)
# zero the 30% smallest-magnitude weights
l1_unstructured(model, name="weight", amount=0.3)
# remove 10 whole output channels (dim 0 of the weight)
ln_structured(model, name="weight", amount=0.1, n=2, dim=0)
print(is_pruned(model)) # True while the mask is active
print(model.weight_mask) # the binary mask, 1 = keepUnstructured methods
Prune |
|
Prune |
|
Prune |
|
Prune several tensors jointly under a single unstructured budget. |
|
Attach the pruning reparameterization to |
|
Zero out the units with the smallest magnitudes. |
|
Zero out a uniformly random subset of the currently unpruned units. |
|
Zero out exactly the units designated by a caller-supplied mask. |
|
Prune nothing and only install the mask reparameterization. |
l1_unstructured() is the standard entry:
prune the amount smallest-|w| entries of one parameter.
custom_from_mask() applies a mask you supply
from domain knowledge. The module-level functions prune one parameter of
one module; global_unstructured() takes an
iterable of (module, name) pairs and a pruning method, and ranks the
entries across all of them together — so two layers compete for the
same global sparsity budget rather than each losing exactly amount%
locally. identity() attaches a
reparameterization that prunes nothing (an all-ones mask), useful as a
placeholder in pipelines that expect one.
Structured methods
Prune |
|
Prune |
|
Zero out the channels with the smallest L``n``-norm. |
|
Zero out entire randomly selected channels of a tensor. |
ln_structured() removes the channels with the
smallest L-norm along dim — n selects the norm order (2 is the common
choice; float('inf'), float('-inf'), 'fro', and 'nuc' are
accepted). Structured pruning zeroes every weight of the
removed channel, so the output dimension shrinks logically — the kind of
sparsity dense kernels and downstream layers can exploit without
specialized sparse support.
Managing the reparameterization
Make the pruning of |
|
Check whether |
|
Check that an absolute pruning count fits the data being pruned. |
|
Convert a pruning amount into an absolute unit count. |
|
Abstract base class for pruning techniques. |
|
Sequence of pruning methods applied iteratively to the same tensor. |
While pruning is active, the parameter (say weight) is moved aside and
the module gains weight_orig (the dense values), weight_mask (the
0/1 mask), and a forward pre-hook that recomputes weight = weight_orig * mask on every call. remove() tears that
scaffolding down and bakes the mask into the stored values, making the
sparsity permanent and the hooks gone;
is_pruned() reports whether any submodule
still carries an active reparameterization.
validate_pruning_amount() and
compute_nparams_to_prune() are the helpers the
methods use to interpret amount and count entries.
BasePruningMethod is the abstract
contract for custom methods: implement compute_mask(t, default_mask)
returning the new mask, and apply it with the classmethod apply(module, name, ...). PruningContainer composes
several methods on one parameter — each pruning call on an
already-pruned parameter unions its mask into the container, so iterative
pruning schedules accumulate.
Where to go next
quantization — the other model-optimization axis: precision instead of sparsity. The two compose (prune, then quantize, or vice versa depending on the target).
nn — the modules being pruned, and
parametrizationsfor the related constraint-style APIs.the main namespace — tensor indexing ops the mask math reduces to.
Help improve this page
Found an error, an unclear step, or a missing example?