# Pruning - tensorplay.ao.pruning

Source: https://www.tensorplay.cn/docs/pruning.html

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 = keep
```

## Unstructured methods

| tensorplay.ao.pruning.l1_unstructured |Prune module[name] by removing the units with the smallest magnitudes. |
| --- | --- |
| tensorplay.ao.pruning.random_unstructured | Prune module[name] by removing a random subset of its units. |
| tensorplay.ao.pruning.custom_from_mask | Prune module[name] with a pre-computed binary mask. |
| tensorplay.ao.pruning.global_unstructured | Prune several tensors jointly under a single unstructured budget. |
| tensorplay.ao.pruning.identity | Attach the pruning reparameterization to module[name] without pruning any unit. |
| tensorplay.ao.pruning.L1Unstructured | Zero out the units with the smallest magnitudes. |
| tensorplay.ao.pruning.RandomUnstructured | Zero out a uniformly random subset of the currently unpruned units. |
| tensorplay.ao.pruning.CustomFromMask | Zero out exactly the units designated by a caller-supplied mask. |
| tensorplay.ao.pruning.Identity | Prune nothing and only install the mask reparameterization. |

[l1_unstructured()](/docs/generated/tensorplay.ao.pruning.l1_unstructured.html#tensorplay.ao.pruning.l1_unstructured) is the standard entry:
prune the amount smallest-|w| entries of one parameter.
[custom_from_mask()](/docs/generated/tensorplay.ao.pruning.custom_from_mask.html#tensorplay.ao.pruning.custom_from_mask) applies a mask you supply
from domain knowledge. The module-level functions prune one parameter of
one module; [global_unstructured()](/docs/generated/tensorplay.ao.pruning.global_unstructured.html#tensorplay.ao.pruning.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()](/docs/generated/tensorplay.ao.pruning.identity.html#tensorplay.ao.pruning.identity) attaches a
reparameterization that prunes nothing (an all-ones mask), useful as a
placeholder in pipelines that expect one.

## Structured methods

| tensorplay.ao.pruning.ln_structured |Prune module[name] by removing the channels with the smallest L``n``-norm along dim. |
| --- | --- |
| tensorplay.ao.pruning.random_structured | Prune module[name] by removing random channels along dim. |
| tensorplay.ao.pruning.LnStructured | Zero out the channels with the smallest L``n``-norm. |
| tensorplay.ao.pruning.RandomStructured | Zero out entire randomly selected channels of a tensor. |

[ln_structured()](/docs/generated/tensorplay.ao.pruning.ln_structured.html#tensorplay.ao.pruning.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

| tensorplay.ao.pruning.remove |Make the pruning of module[name] permanent and drop the reparameterization. |
| --- | --- |
| tensorplay.ao.pruning.is_pruned | Check whether module carries an active pruning reparameterization. |
| tensorplay.ao.pruning.validate_pruning_amount | Check that an absolute pruning count fits the data being pruned. |
| tensorplay.ao.pruning.compute_nparams_to_prune | Convert a pruning amount into an absolute unit count. |
| tensorplay.ao.pruning.BasePruningMethod | Abstract base class for pruning techniques. |
| tensorplay.ao.pruning.PruningContainer | 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()](/docs/generated/tensorplay.ao.pruning.remove.html#tensorplay.ao.pruning.remove) tears that
scaffolding down and bakes the mask into the stored values, making the
sparsity permanent and the hooks gone;
[is_pruned()](/docs/generated/tensorplay.ao.pruning.is_pruned.html#tensorplay.ao.pruning.is_pruned) reports whether any submodule
still carries an active reparameterization.
[validate_pruning_amount()](/docs/generated/tensorplay.ao.pruning.validate_pruning_amount.html#tensorplay.ao.pruning.validate_pruning_amount) and
[compute_nparams_to_prune()](/docs/generated/tensorplay.ao.pruning.compute_nparams_to_prune.html#tensorplay.ao.pruning.compute_nparams_to_prune) are the helpers the
methods use to interpret amount and count entries.

[BasePruningMethod](/docs/generated/tensorplay.ao.pruning.BasePruningMethod.html#tensorplay.ao.pruning.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](/docs/generated/tensorplay.ao.pruning.PruningContainer.html#tensorplay.ao.pruning.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](/docs/quantization.html) — the other model-optimization axis: precision instead of sparsity. The two compose (prune, then quantize, or vice versa depending on the target).

- [nn](/docs/nn.html) — the modules being pruned, and parametrizations for the related constraint-style APIs.

- [the main namespace](/docs/tensorplay.html) — tensor indexing ops the mask math reduces to.