tensorplay.ao.pruning API

Functions 11

compute_nparams_to_prune

tensorplay.ao.pruning.compute_nparams_to_prune(amount: int | float, tensor_size: int) → int[source]

Convert a pruning amount into an absolute unit count.

An integer amount is returned unchanged; a float amount is read as a fraction of tensor_size and rounded to the nearest integer.

Parameters:

amount – pruning quantity (absolute int, or fractional float in [0, 1]).
tensor_size – number of units available for pruning.

Returns:

The number of units to prune.

custom_from_mask

functionFull reference ↗

tensorplay.ao.pruning.custom_from_mask(module: Module, name: str, mask: Tensor) → Module[source]

Prune module[name] with a pre-computed binary mask.

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

adding a named buffer called name + '_mask' holding mask;
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.
mask – binary mask to be applied to the parameter.

Returns:

The modified (i.e. pruned) module.

Examples

>>> # xdoctest: +SKIP
>>> m = custom_from_mask(
...     nn.Linear(5, 3), name="bias", mask=tensorplay.tensor([0, 1, 0])
... )
>>> print(m.bias_mask)
tensor([0., 1., 0.])

global_unstructured

functionFull reference ↗

tensorplay.ao.pruning.global_unstructured(parameters: Iterable[tuple[Module, str]], pruning_method: type[BasePruningMethod], importance_scores: dict[tuple[Module, str], Tensor] | None = None, **kwargs: Any) → None[source]

Prune several tensors jointly under a single unstructured budget.

Aggregates the importance scores of every listed parameter into one vector, computes a single mask under the shared amount budget, and slices that mask back onto each parameter. Modifies the modules in place by:

adding a named buffer called name + '_mask' for every listed parameter;
replacing each parameter name by its masked version, while the original (unmasked) values are stored in a new parameter named name + '_orig'.

Parameters:

parameters – iterable of (module, name) tuples identifying the parameters to prune globally, i.e. by aggregating all values before deciding which units to remove.
pruning_method – a pruning method class from this package (or a user-defined subclass of BasePruningMethod) whose PRUNING_TYPE is 'unstructured'.
importance_scores – mapping from (module, name) tuples to the corresponding importance-scores tensor (same shape as the parameter). Parameters absent from the mapping use their own values as importance scores.
kwargs – keyword arguments forwarded to pruning_method, typically amount: the quantity of units to prune across all listed parameters (a float fraction in [0, 1] or an absolute int).

Raises:

TypeError – if parameters is not an iterable, if importance_scores is not a dict, or if the PRUNING_TYPE of pruning_method is not 'unstructured'.

Note

Global pruning is restricted to unstructured methods: a structured norm is only comparable across channels of equal size, which cannot be guaranteed across heterogeneous parameters.

Examples

>>> # xdoctest: +SKIP
>>> net = nn.Sequential(nn.Linear(10, 4), nn.Linear(4, 1))
>>> parameters_to_prune = (
...     (net[0], "weight"),
...     (net[1], "weight"),
... )
>>> global_unstructured(
...     parameters_to_prune,
...     pruning_method=L1Unstructured,
...     amount=10,
... )

identity

functionFull reference ↗

tensorplay.ao.pruning.identity(module: Module, name: str) → Module[source]

Attach the pruning reparameterization to module[name] without pruning any unit.

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

adding a named buffer called name + '_mask' holding the binary mask applied to the parameter name;
replacing the parameter name by its masked version, while the original (unmasked) values are stored in a new parameter named name + '_orig'.

Note

The mask is a tensor of ones.

Parameters:

module – module containing the tensor to prune.
name – parameter name within module on which pruning acts.

Returns:

The modified (i.e. pruned) module.

Examples

>>> # xdoctest: +SKIP
>>> m = identity(nn.Linear(2, 3), "bias")
>>> print(m.bias_mask)
tensor([1., 1., 1.])

is_pruned

functionFull reference ↗

tensorplay.ao.pruning.is_pruned(module: Module) → bool[source]

Check whether module carries an active pruning reparameterization.

Scans every submodule for forward pre-hooks that are instances of BasePruningMethod.

Parameters:: module – module that is either pruned or unpruned.
Returns:: True when at least one submodule is pruned, False otherwise.

Examples

>>> # xdoctest: +SKIP
>>> m = nn.Linear(5, 7)
>>> print(is_pruned(m))
False
>>> random_unstructured(m, name="weight", amount=0.2)
>>> print(is_pruned(m))
True

l1_unstructured

functionFull reference ↗

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

Prune module[name] by removing the units with the smallest magnitudes.

Removes amount (currently unpruned) units ranked by absolute value. Modifies the module in place (and also returns it) by:

adding a named buffer called name + '_mask' holding the binary mask applied to the parameter name;
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 units to prune. A float in [0, 1] denotes the fraction of units to prune; an int denotes the absolute number of units to prune.
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 = l1_unstructured(nn.Linear(2, 3), "weight", amount=0.2)
>>> list(m.state_dict().keys())
['bias', 'weight_orig', 'weight_mask']

ln_structured

functionFull reference ↗

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:

adding a named buffer called name + '_mask' holding the binary mask applied to the parameter name;
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")
... )

random_structured

functionFull reference ↗

tensorplay.ao.pruning.random_structured(module: Module, name: str, amount: int | float, dim: int) → Module[source]

Prune module[name] by removing random channels along dim.

Removes amount (currently unpruned) channels chosen uniformly at random. Modifies the module in place (and also returns it) by:

adding a named buffer called name + '_mask' holding the binary mask applied to the parameter name;
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.
dim – axis along which channels are defined.

Returns:

The modified (i.e. pruned) module.

Examples

>>> # xdoctest: +SKIP
>>> m = random_structured(nn.Linear(5, 3), "weight", amount=3, dim=1)
>>> columns_pruned = int(sum(tensorplay.sum(m.weight, dim=0) == 0))
>>> print(columns_pruned)
3

random_unstructured

functionFull reference ↗

tensorplay.ao.pruning.random_unstructured(module: Module, name: str, amount: int | float) → Module[source]

Prune module[name] by removing a random subset of its units.

Removes amount (currently unpruned) units chosen uniformly at random. Modifies the module in place (and also returns it) by:

adding a named buffer called name + '_mask' holding the binary mask applied to the parameter name;
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 units to prune. A float in [0, 1] denotes the fraction of units to prune; an int denotes the absolute number of units to prune.

Returns:

The modified (i.e. pruned) module.

Examples

>>> # xdoctest: +SKIP
>>> m = random_unstructured(nn.Linear(2, 3), "weight", amount=1)
>>> int(tensorplay.sum(m.weight_mask == 0))
1

remove

functionFull reference ↗

tensorplay.ao.pruning.remove(module: Module, name: str) → Module[source]

Make the pruning of module[name] permanent and drop the reparameterization.

The pruned parameter name remains permanently pruned, the parameter name + '_orig' is removed from the parameter list, and the buffer name + '_mask' is removed from the buffers. The pruning hook is also detached from the module.

Note

Pruning itself is NOT undone or reversed!

Parameters:

module – module containing the pruned tensor.
name – parameter name within module whose pruning is removed.

Returns:

The modified module.

Raises:

ValueError – if name is not currently pruned on module.

validate_pruning_amount

functionFull reference ↗

tensorplay.ao.pruning.validate_pruning_amount(amount: int, tensor_size: int) → None[source]

Check that an absolute pruning count fits the data being pruned.

Parameters:

amount – absolute number of units to prune.
tensor_size – number of units available for pruning.

Raises:

ValueError – if amount exceeds tensor_size.

Classes 8

BasePruningMethod

classFull reference ↗

class tensorplay.ao.pruning.BasePruningMethod[source]

Abstract base class for pruning techniques.

Subclasses must override compute_mask() and declare a PRUNING_TYPE class attribute (one of 'unstructured', 'structured' or 'global'). The classmethod apply() installs the reparameterization on a module and registers an instance of the subclass as a forward pre-hook.

classmethod apply(module: Module, name: str, *args: Any, importance_scores: Tensor | None = None, **kwargs: Any) → BasePruningMethod[source]

Install the pruning reparameterization for module[name].

Moves the original parameter to name + '_orig', registers the computed mask as the buffer name + '_mask', stores the masked values under name and adds a forward pre-hook that re-applies the mask on every forward pass. If the tensor is already pruned by another method, the new method is composed with the existing one through a PruningContainer.

Parameters:

module – module containing the tensor to prune.
name – parameter name within module on which pruning acts.
args – positional arguments forwarded to the subclass constructor.
importance_scores – tensor of importance scores with the same shape as module[name]; each entry ranks the corresponding element of the parameter. When unspecified, the parameter itself is used as its own importance scores.
kwargs – keyword arguments forwarded to the subclass constructor.

Returns:

The pruning method (or container of methods) now attached to the module.

apply_mask(module: Module) → Tensor[source]

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.

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

Compute the pruning mask for the input tensor t.

Starting from default_mask (a mask of ones when t has never been pruned), derive the new mask according to the recipe of the concrete method. Entries already zeroed by default_mask must stay zero.

Parameters:

t – tensor whose entries or channels are ranked for pruning, typically the importance scores of the parameter.
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.

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

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[source]

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!

CustomFromMask

classFull reference ↗

class tensorplay.ao.pruning.CustomFromMask(mask: Tensor)[source]

Zero out exactly the units designated by a caller-supplied mask.

Parameters:: mask – binary mask whose zeros mark the units to prune.

classmethod apply(module: Module, name: str, mask: Tensor) → BasePruningMethod[source]

Install a user-provided mask for module[name].

Parameters:

module – module containing the tensor to prune.
name – parameter name within module on which pruning acts.
mask – binary mask to be applied to the parameter.

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.

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.

Note

Pruning itself is NOT undone or reversed!

Identity

classFull reference ↗

class tensorplay.ao.pruning.Identity[source]

Prune nothing and only install the mask reparameterization.

The generated mask is a tensor of ones, which is useful to prepare a module for later iterative pruning without removing any unit yet.

classmethod apply(module: Module, name: str) → BasePruningMethod[source]

Install the identity (all-ones mask) reparameterization.

Parameters:

module – module containing the tensor to reparameterize.
name – parameter name within module on which pruning acts.

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.

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.

Note

Pruning itself is NOT undone or reversed!

L1Unstructured

classFull reference ↗

class tensorplay.ao.pruning.L1Unstructured(amount: int | float)[source]

Zero out the units with the smallest magnitudes.

Ranks all currently unpruned units by absolute value and removes the amount smallest ones.

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

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

Install magnitude-based unstructured pruning for module[name].

Parameters:

module – module containing the tensor to prune.
name – parameter name within module on which pruning acts.
amount – quantity of units to prune. A float in [0, 1] denotes the fraction of units to prune; an int denotes the absolute number of units to prune.
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.

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.

Note

Pruning itself is NOT undone or reversed!

LnStructured

classFull reference ↗

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.

Note

Pruning itself is NOT undone or reversed!

PruningContainer

classFull reference ↗

class tensorplay.ao.pruning.PruningContainer(*args: BasePruningMethod)[source]

Sequence of pruning methods applied iteratively to the same tensor.

Tracks the order in which methods were added and combines successive pruning calls: each new method only ranks the entries or channels that the previous masks left unpruned.

Accepts as argument an instance of a BasePruningMethod or an iterable of them.

add_pruning_method(method: BasePruningMethod | None) → None[source]

Add a child pruning method to the container.

Parameters:

method – child pruning method to be added to the container.

Raises:

TypeError – if method is neither None nor a BasePruningMethod instance.
ValueError – if method acts on a different tensor name than the methods already held by the container.

classmethod apply(module: Module, name: str, *args: Any, importance_scores: Tensor | None = None, **kwargs: Any) → BasePruningMethod

Install the pruning reparameterization for module[name].

Parameters:

module – module containing the tensor to prune.
name – parameter name within module on which pruning acts.
args – positional arguments forwarded to the subclass constructor.
importance_scores – tensor of importance scores with the same shape as module[name]; each entry ranks the corresponding element of the parameter. When unspecified, the parameter itself is used as its own importance scores.
kwargs – keyword arguments forwarded to the subclass constructor.

Returns:

The pruning method (or container of methods) now attached to the module.

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]

Apply the latest method and merge its mask into default_mask.

The new partial mask is computed on the entries or channels that default_mask has not zeroed out. Which portion of t the new mask is derived from depends on the PRUNING_TYPE of the last method:

'unstructured': the mask is computed from the flattened list of entries not yet masked;
'structured': the mask is computed from the channels that still hold at least one unmasked entry;
'global': the mask is computed across all entries.

Parameters:

t – tensor representing the parameter to prune (same shape as default_mask).
default_mask – mask accumulated from previous pruning iterations.

Returns:

The mask combining the effects of default_mask and of the latest method, with the same shape as default_mask and 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.

Note

Pruning itself is NOT undone or reversed!

RandomStructured

classFull reference ↗

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

Zero out entire randomly selected channels of a tensor.

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.
dim – axis along which channels are defined. Default: -1.

classmethod apply(module: Module, name: str, amount: int | float, dim: int = -1) → BasePruningMethod[source]

Install random 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.
dim – axis along which channels are defined. Default: -1.

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 drawing channels at random.

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.

Note

Pruning itself is NOT undone or reversed!

RandomUnstructured

classFull reference ↗

class tensorplay.ao.pruning.RandomUnstructured(amount: int | float)[source]

Zero out a uniformly random subset of the currently unpruned units.

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

classmethod apply(module: Module, name: str, amount: int | float) → BasePruningMethod[source]

Install random unstructured pruning for module[name].

Parameters:

module – module containing the tensor to prune.
name – parameter name within module on which pruning acts.
amount – quantity of units to prune. A float in [0, 1] denotes the fraction of units to prune; an int denotes the absolute number of units to prune.

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.

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.

Note

Pruning itself is NOT undone or reversed!

Functions 11

compute_nparams_to_prune

custom_from_mask

global_unstructured

identity

is_pruned

l1_unstructured

ln_structured

random_structured

random_unstructured

remove

validate_pruning_amount

Classes 8

BasePruningMethod

CustomFromMask

Identity

L1Unstructured

LnStructured

PruningContainer

RandomStructured

RandomUnstructured

Help improve this page

On this page