tensorplay.ao.pruning.random_unstructured

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:

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 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