unstructured

unstructured

Global L1 unstructured weight pruning.

UnstructuredPruner

UnstructuredPruner(amount, layer_types=None)

Bases: BasePruner

Global L1 unstructured pruning across all target layers.

Initialize with pruning amount and target layer types.

Parameters:

Name	Type	Description	Default
amount	float	Fraction of weights to zero, must be in (0, 1).	required
layer_types	list[str] | None	PrunableLayerType values to target. Defaults to convolution and linear layers (normalization scales and embedding tables are usually not good pruning targets).	None

Raises:

Type	Description
ValueError	If amount is not in the open interval (0, 1).

Source code in src/versatil/post_training_compression/pruning/unstructured.py

def __init__(
    self,
    amount: float,
    layer_types: list[str] | None = None,
) -> None:
    """Initialize with pruning amount and target layer types.

    Args:
        amount: Fraction of weights to zero, must be in (0, 1).
        layer_types: PrunableLayerType values to target. Defaults to
            convolution and linear layers (normalization scales and
            embedding tables are usually not good pruning targets).

    Raises:
        ValueError: If amount is not in the open interval (0, 1).
    """
    if amount <= 0.0 or amount >= 1.0:
        raise ValueError(f"Pruning amount must be in (0, 1), got {amount}")
    self._amount = amount
    if layer_types is None:
        layer_types = [
            PrunableLayerType.CONV1D.value,
            PrunableLayerType.CONV2D.value,
            PrunableLayerType.LINEAR.value,
        ]
    self._layer_types = tuple(
        PrunableLayerType(name).to_module_type() for name in layer_types
    )

amount property

amount

Fraction of weights to prune.

layer_types property

layer_types

Module types targeted for pruning. None means all.

prune

prune(module)

Apply global L1 unstructured pruning.

Identifies all target layers, applies global L1 unstructured pruning, then removes the pruning reparametrization to make weights permanent.

Parameters:

Name	Type	Description	Default
module	Module	Neural network module to prune.	required

Returns:

Type	Description
tuple[int, int]	Tuple of (total_parameters, zero_parameters).

Source code in src/versatil/post_training_compression/pruning/unstructured.py

def prune(self, module: nn.Module) -> tuple[int, int]:
    """Apply global L1 unstructured pruning.

    Identifies all target layers, applies global L1 unstructured pruning,
    then removes the pruning reparametrization to make weights permanent.

    Args:
        module: Neural network module to prune.

    Returns:
        Tuple of (total_parameters, zero_parameters).
    """
    parameters_to_prune = [
        (child_module, PruningTargetAttribute.WEIGHT.value)
        for child_module in module.modules()
        if isinstance(
            getattr(child_module, PruningTargetAttribute.WEIGHT.value, None),
            nn.Parameter,
        )
        and isinstance(child_module, self._layer_types)
    ]
    if not parameters_to_prune:
        raise ValueError(
            "Unstructured pruning selected no modules; the target module "
            f"contains no {[t.__name__ for t in self._layer_types]} layers."
        )
    prune.global_unstructured(
        parameters_to_prune,
        pruning_method=prune.L1Unstructured,  # type: ignore[arg-type]
        amount=self._amount,
    )
    for child_module, name in parameters_to_prune:
        prune.remove(child_module, name)
    return self.compute_sparsity(module)