LnStructured

class tensorplay.ao.pruning.LnStructured(amount: int | float, n: int | float | str, dim: int = -1)[source]

Zero out the channels with the smallest L``n``-norm.

Reduces the tensor to one L``n``-norm per channel along dim and keeps the amount-largest channels, zeroing the rest.

Parameters:

• 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() (e.g. a positive number, float('inf'), float('-inf'), 'fro', 'nuc').

• dim – axis along which channels are defined. Default: -1.

classmethod apply(module: Module, name: str, amount: int | float, n: int | float | str, dim: int, importance_scores: Tensor | None = None) → BasePruningMethod[source]

Install norm-based structured pruning for module[name].

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.

apply_mask(module: Module) → Tensor

Return the pruned version of the tensor held by module.

Fetches the mask and the original values from the module and returns their elementwise product.

Parameters:

module – module holding the pruned tensor.

Returns:

The product of the mask and the original values.

compute_mask(t: Tensor, default_mask: Tensor) → Tensor[source]

Compute a channel mask for t by per-channel L``n``-norm.

Parameters:

• t – tensor whose channels are pruned.

• default_mask – mask accumulated from previous pruning iterations; must be respected by the new mask. Same shape as t.

Returns:

The mask to apply to t, with the same shape as t.

Raises:

IndexError – if self.dim is not a valid axis of t.

prune(t: Tensor, default_mask: Tensor | None = None, importance_scores: Tensor | None = None) → Tensor

Return a pruned copy of the input tensor t.

Applies the rule implemented by compute_mask() without any module-side reparameterization.

Parameters:

• t – tensor to prune (same shape as default_mask).

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

• default_mask – mask from a previous pruning iteration, if any. Pruning must respect the entries it already zeroes. When unspecified, a mask of ones is used.

Returns:

The pruned version of t.

remove(module: Module) → None

Make the current pruning of module permanent.

The pruned values remain pruned: the product of mask and original values is written back into the parameter name, and the auxiliary parameter name + '_orig' and buffer name + '_mask' are dropped.

Note

Pruning itself is NOT undone or reversed!