tensorplay.ao.pruning.ln_structured

tensorplay.ao.pruning.ln_structured(module: Module, name: str, amount: int | float, n: int | float | str, dim: int, importance_scores: Tensor | None = None) → Module[source]

Prune module[name] by removing the channels with the smallest L``n``-norm along dim.

Modifies the module in place (and also returns it) by:

1. adding a named buffer called name + '_mask' holding the binary mask applied to the parameter name;

2. replacing the parameter name by its masked version, while the original (unmasked) values are stored in a new parameter named name + '_orig'.

Parameters:

• module – module containing the tensor to prune.

• name – parameter name within module on which pruning acts.

• amount – quantity of channels to prune. A float in [0, 1] denotes the fraction of channels to prune; an int denotes the absolute number of channels to prune.

• n – norm order; accepts the orders valid for tensorplay.linalg.vector_norm() and tensorplay.linalg.matrix_norm().

• dim – axis along which channels are defined.

• importance_scores – tensor of importance scores with the same shape as the parameter; each entry ranks the corresponding element of the parameter. When unspecified, the parameter itself is used.

Returns:

The modified (i.e. pruned) module.

Examples

>>> # xdoctest: +SKIP
>>> m = ln_structured(
...     nn.Conv2d(5, 3, 2), "weight", amount=0.3, dim=1, n=float("-inf")
... )