tensorplay.ao.pruning.identity

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:

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

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.])