Skip to content

structured

structured

Per-layer structured weight pruning using Lp-norm channel ranking.

StructuredPruner

StructuredPruner(amount, norm_order=2, dimension=0, layer_types=None)

Bases: BasePruner

Per-layer structured pruning along a specified dimension.

Ranks channels by their Lp-norm magnitude and zeros the lowest-ranked fraction. The norm_order parameter specifies p in the Lp-norm (e.g., 1 for L1-norm, 2 for L2-norm).

Initialize with pruning parameters.

Parameters:

Name Type Description Default
amount float

Fraction of channels to zero per layer, must be in (0, 1).

required
norm_order int

The p in Lp-norm used to rank channels (e.g., 1 for L1, 2 for L2).

2
dimension int

Weight tensor dimension along which to prune.

0
layer_types list[str] | None

PrunableLayerType values to target. Defaults to Conv1d, Conv2d, and Linear.

None

Raises:

Type Description
ValueError

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

Source code in src/versatil/post_training_compression/pruning/structured.py
def __init__(
    self,
    amount: float,
    norm_order: int = 2,
    dimension: int = 0,
    layer_types: list[str] | None = None,
) -> None:
    """Initialize with pruning parameters.

    Args:
        amount: Fraction of channels to zero per layer, must be in (0, 1).
        norm_order: The p in Lp-norm used to rank channels
            (e.g., 1 for L1, 2 for L2).
        dimension: Weight tensor dimension along which to prune.
        layer_types: PrunableLayerType values to target. Defaults
            to Conv1d, Conv2d, and Linear.

    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}")
    if layer_types is None:
        layer_types = [
            PrunableLayerType.CONV1D.value,
            PrunableLayerType.CONV2D.value,
            PrunableLayerType.LINEAR.value,
        ]
    self._amount = amount
    self._norm_order = norm_order
    self._dimension = dimension
    self._layer_types = tuple(
        PrunableLayerType(name).to_module_type() for name in layer_types
    )

amount property

amount

Fraction of channels to prune per layer.

norm_order property

norm_order

Ln norm order for channel ranking.

dimension property

dimension

Weight tensor dimension along which pruning is applied.

layer_types property

layer_types

Module types targeted for pruning.

prune

prune(module)

Apply per-layer Ln structured pruning.

Iterates over all target layers and applies Ln structured pruning individually, then removes the pruning reparametrization.

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

Raises:

Type Description
ValueError

If the target module contains no prunable layers.

Source code in src/versatil/post_training_compression/pruning/structured.py
def prune(self, module: nn.Module) -> tuple[int, int]:
    """Apply per-layer Ln structured pruning.

    Iterates over all target layers and applies Ln structured pruning
    individually, then removes the pruning reparametrization.

    Args:
        module: Neural network module to prune.

    Returns:
        Tuple of (total_parameters, zero_parameters).

    Raises:
        ValueError: If the target module contains no prunable layers.
    """
    target_layers = [
        child_module
        for child_module in module.modules()
        if isinstance(child_module, self._layer_types)
    ]
    if not target_layers:
        raise ValueError(
            "Structured pruning selected no modules; the target module "
            f"contains no {[t.__name__ for t in self._layer_types]} layers."
        )
    for child_module in target_layers:
        prune.ln_structured(
            child_module,
            name=PruningTargetAttribute.WEIGHT.value,
            amount=self._amount,
            n=self._norm_order,
            dim=self._dimension,
        )
        prune.remove(child_module, PruningTargetAttribute.WEIGHT.value)
    return self.compute_sparsity(module)