pruning

shrinkai.compression.pruning

Pruning: inducing sparsity or physically shrinking a model.

Two complementary tools:

• Pruner (configured via PruningConfig) applies unstructured or structured masking through torch.nn.utils.prune. It never changes tensor shapes — even "structured" pruning here only zeroes out channels, so it does not by itself reduce parameter count, disk size, or latency on standard hardware.
• ChannelPruner physically removes pruned channels from Conv2d/Linear layers (and their dependent BatchNorm), using torch.fx to trace the model's dataflow graph and safely propagate the shrink downstream. This is what actually reduces parameters, size, latency, and FLOPs — but only on simple, non-branching topologies (see its own docstring for the exact scope).

Modules:

Name	Description
channel_pruner	Physical (dependency-aware) structured channel pruning.
pruning

Classes:

Name	Description
ChannelPruner	Physically removes pruned output channels from Conv2d/Linear layers.
Pruner	Orchestrates the pruning of PyTorch models to induce sparsity.
PruningConfig	Configuration parameters for model pruning.

Classes

ChannelPruner

Physically removes pruned output channels from Conv2d/Linear layers.

See the module docstring for the exact topologies this supports and rejects.

Equation

Output units (rows of the weight tensor) are ranked by their \(L_2\) norm, following the filter-pruning criterion of Li et al. (2017) (who originally used the \(L_1\) norm; this implementation uses \(L_2\), matching Pruner's structured method):

\[\|W_j\|_2 = \left(\sum_i w_{j,i}^2\right)^{1/2}\]

The amount fraction of channels with the smallest norm are removed, physically, unlike Pruner, whose masking-based criterion is identical but only zeroes the weights without changing tensor shapes.

Parameters:

Name	Type	Description	Default
amount	float	Fraction of output channels/neurons to remove from each targeted layer, ranked by L2-norm (lowest-norm channels removed first). Must be in (0.0, 1.0). Defaults to 0.3.	0.3

Methods:

Name	Description
apply	Physically prunes the given layers, propagating the shrink downstream.
benchmark	Compares the original model against the physically pruned one.

Source code in src/shrinkai/compression/pruning/channel_pruner.py

class ChannelPruner:
    r"""Physically removes pruned output channels from `Conv2d`/`Linear` layers.

    See the module docstring for the exact topologies this supports and rejects.

    Equation:
        Output units (rows of the weight tensor) are ranked by their $L_2$ norm,
        following the filter-pruning criterion of Li et al. (2017) (who originally
        used the $L_1$ norm; this implementation uses $L_2$, matching `Pruner`'s
        structured method):

        $$\|W_j\|_2 = \left(\sum_i w_{j,i}^2\right)^{1/2}$$

        The `amount` fraction of channels with the smallest norm are removed,
        physically, unlike `Pruner`, whose masking-based criterion is identical
        but only zeroes the weights without changing tensor shapes.

    Args:
        amount: Fraction of output channels/neurons to remove from each targeted
            layer, ranked by L2-norm (lowest-norm channels removed first). Must be
            in (0.0, 1.0). Defaults to 0.3.
    """  # noqa: E501

    def __init__(self, amount: float = 0.3) -> None:
        if not (0.0 < amount < 1.0):
            raise ValueError(f"amount must be between 0.0 and 1.0, got {amount}")
        self.amount = amount

    def apply(
        self,
        model: nn.Module,
        target_layers: list[str],
        sample_input: torch.Tensor,
    ) -> nn.Module:
        """Physically prunes the given layers, propagating the shrink downstream.

        Args:
            model: The model to prune. Mutated in place (its pruned submodules are
                replaced) and also returned for convenience.
            target_layers: Names (as in `model.named_modules()`) of `Conv2d`/`Linear`
                layers whose output channels/neurons should be pruned.
            sample_input: A representative input tensor, used to trace the model's
                dataflow graph and determine each node's actual tensor shape. Not
                used for any weight-affecting computation.

        Returns:
            nn.Module: The same `model`, with the targeted layers (and any
            dependent BatchNorm / downstream layer) physically shrunk.

        Raises:
            ValueError: If a target layer does not exist, is not a `Conv2d`/`Linear`
                with `groups == 1`, or its output cannot be safely traced to a
                single downstream layer (branching, unsupported op, or an unsafe
                flatten across non-unit spatial dimensions).
        """
        traced = torch.fx.symbolic_trace(model)
        ShapeProp(traced).propagate(sample_input)
        node_by_target = {n.target: n for n in traced.graph.nodes if n.op == "call_module"}
        # Used only to type-check nodes while walking the graph (Conv2d vs BatchNorm
        # vs activation, ...). Structurally identical to `model`'s own module tree,
        # so it stays valid for type-checking even after `model` is mutated below.
        traced_modules = dict(traced.named_modules())

        for layer_name in target_layers:
            self._prune_one_layer(model, layer_name, node_by_target, traced_modules)

        return model

    def _prune_one_layer(
        self,
        model: nn.Module,
        layer_name: str,
        node_by_target: dict,
        traced_modules: dict,
    ) -> None:
        modules = dict(model.named_modules())
        if layer_name not in modules:
            raise ValueError(f"Layer '{layer_name}' not found in the model.")
        if layer_name not in node_by_target:
            raise ValueError(
                f"Layer '{layer_name}' was not found as a traced `call_module` node. "
                "It may be unused in the forward pass, or called in a way "
                "`torch.fx.symbolic_trace` could not capture."
            )

        layer = modules[layer_name]
        if not isinstance(layer, _PRUNABLE_TYPES):
            raise ValueError(
                f"Layer '{layer_name}' is a {type(layer).__name__}, but ChannelPruner "
                "only supports Conv2d and Linear layers."
            )
        if isinstance(layer, nn.Conv2d) and layer.groups != 1:
            raise ValueError(
                f"Layer '{layer_name}' is a grouped/depthwise convolution (groups="
                f"{layer.groups}), which ChannelPruner does not support."
            )

        bn_name, downstream_name = self._walk_dependents(node_by_target[layer_name], traced_modules)

        keep_indices = _rank_keep_indices(layer.weight, self.amount)
        logger.info(
            "ChannelPruner: pruning '%s' from %d to %d output channels.",
            layer_name,
            layer.weight.shape[0],
            len(keep_indices),
        )
        _set_module_by_name(model, layer_name, _shrink_output(layer, keep_indices))

        if bn_name is not None:
            bn_layer = dict(model.named_modules())[bn_name]
            _set_module_by_name(model, bn_name, _shrink_batchnorm(bn_layer, keep_indices))

        if downstream_name is not None:
            downstream_layer = dict(model.named_modules())[downstream_name]
            if isinstance(downstream_layer, nn.Conv2d) and downstream_layer.groups != 1:
                raise ValueError(
                    f"Downstream layer '{downstream_name}' of '{layer_name}' is a "
                    f"grouped/depthwise convolution (groups={downstream_layer.groups}), "
                    "which ChannelPruner does not support."
                )
            _set_module_by_name(
                model, downstream_name, _shrink_input(downstream_layer, keep_indices)
            )

    def _walk_dependents(self, node, traced_modules: dict) -> tuple[str | None, str | None]:
        """Walks forward from `node` to find a co-prunable BatchNorm and the next
        `Conv2d`/`Linear` layer whose input channels must shrink to match.

        Args:
            node: The `call_module` FX node of the layer being pruned.
            traced_modules: `dict(traced_graph_module.named_modules())`, used to
                type-check the modules encountered while walking forward.

        Returns:
            tuple[str | None, str | None]: (batchnorm_layer_name, downstream_layer_name),
            either of which may be None (no BatchNorm in the path / target is the
            model's final layer).
        """
        current = node
        bn_name: str | None = None

        while True:
            if len(current.users) != 1:
                raise ValueError(
                    f"Layer '{node.target}' output is consumed by "
                    f"{len(current.users)} node(s) (expected exactly 1). Branching "
                    "topologies (skip connections, concatenation, ...) are not "
                    "supported by ChannelPruner; use `Pruner` instead."
                )

            next_node = next(iter(current.users))

            if next_node.op == "output":
                return bn_name, None

            if next_node.op == "call_module":
                next_module = traced_modules[next_node.target]
                if isinstance(next_module, _BATCHNORM_TYPES):
                    bn_name = next_node.target
                    current = next_node
                    continue
                if isinstance(next_module, _PRUNABLE_TYPES):
                    return bn_name, next_node.target
                if isinstance(next_module, _CHANNEL_PRESERVING_MODULE_TYPES):
                    current = next_node
                    continue
                raise ValueError(
                    f"Layer '{node.target}' output passes through unsupported "
                    f"module '{next_node.target}' ({type(next_module).__name__}) "
                    "before reaching a Conv2d/Linear layer. ChannelPruner cannot "
                    "guarantee this is safe to prune through."
                )

            if next_node.op in ("call_function", "call_method"):
                if next_node.target in _RESHAPE_TARGETS:
                    shape = current.meta["tensor_meta"].shape
                    if len(shape) > 2 and any(dim != 1 for dim in shape[2:]):
                        raise ValueError(
                            f"Layer '{node.target}' output is reshaped while its "
                            f"spatial dimensions are {tuple(shape[2:])}, not all 1. "
                            "Pruning through a flatten/view/reshape is only safe "
                            "once spatial size has been reduced to 1x1 (e.g. via "
                            "AdaptiveAvgPool2d(1))."
                        )
                    current = next_node
                    continue
                if next_node.target in _CHANNEL_PRESERVING_FUNCTION_TARGETS:
                    current = next_node
                    continue

            raise ValueError(
                f"Layer '{node.target}' output passes through unsupported "
                f"{next_node.op} '{next_node.target}' before reaching a "
                "Conv2d/Linear layer. ChannelPruner cannot guarantee this is safe "
                "to prune through; use `Pruner` instead for this topology."
            )

    def benchmark(
        self,
        original_model: nn.Module,
        pruned_model: nn.Module,
        sample_input: torch.Tensor,
        original_name: str = "Original (Dense)",
        pruned_name: str = "Pruned (Physically Shrunk)",
        val_dataloader: DataLoader | None = None,
        device: str | torch.device = "cpu",
        compute_flops: bool = False,
    ) -> BenchmarkReport:
        """Compares the original model against the physically pruned one.

        Unlike `Pruner.benchmark`, no `.finalize()`-style step is needed first:
        `pruned_model` (as returned by `.apply()`) already has fewer parameters,
        so real gains in size, latency, and FLOPs are expected here, not just
        theoretical sparsity.

        Args:
            original_model: The original, unpruned model.
            pruned_model: The model returned by `.apply()`.
            sample_input: A dummy tensor for latency measurement.
            original_name: Display label for the original model.
            pruned_name: Display label for the pruned model.
            val_dataloader: Optional dataloader for accuracy comparison.
            device: Device for the benchmark.
            compute_flops: If True, also reports FLOPs per sample for both
                models. Safe to enable here (unlike for `Pruner`/`Quantizer`),
                since physical channel pruning keeps standard Conv2d/Linear ops
                that FLOPs counting handles accurately. Defaults to False.

        Returns:
            BenchmarkReport: Structured benchmark report ready for `.show()`.
        """
        original_acc: float | None = None
        pruned_acc: float | None = None

        if val_dataloader is not None:
            original_acc = compute_accuracy(original_model, val_dataloader, device)
            pruned_acc = compute_accuracy(pruned_model, val_dataloader, device)

        return Profiler.compare(
            teacher=original_model,
            student=pruned_model,
            sample_input=sample_input,
            device=device,
            teacher_name=original_name,
            student_name=pruned_name,
            teacher_acc=original_acc,
            student_acc=pruned_acc,
            compute_flops=compute_flops,
        )

Methods:

apply

apply(
    model: Module,
    target_layers: list[str],
    sample_input: Tensor,
) -> nn.Module

Physically prunes the given layers, propagating the shrink downstream.

Parameters:

Name	Type	Description	Default
model	Module	The model to prune. Mutated in place (its pruned submodules are replaced) and also returned for convenience.	required
target_layers	list[str]	Names (as in model.named_modules()) of Conv2d/Linear layers whose output channels/neurons should be pruned.	required
sample_input	Tensor	A representative input tensor, used to trace the model's dataflow graph and determine each node's actual tensor shape. Not used for any weight-affecting computation.	required

Returns:

Type	Description
Module	nn.Module: The same model, with the targeted layers (and any
Module	dependent BatchNorm / downstream layer) physically shrunk.

Raises:

Type	Description
ValueError	If a target layer does not exist, is not a Conv2d/Linear with groups == 1, or its output cannot be safely traced to a single downstream layer (branching, unsupported op, or an unsafe flatten across non-unit spatial dimensions).

Source code in src/shrinkai/compression/pruning/channel_pruner.py

def apply(
    self,
    model: nn.Module,
    target_layers: list[str],
    sample_input: torch.Tensor,
) -> nn.Module:
    """Physically prunes the given layers, propagating the shrink downstream.

    Args:
        model: The model to prune. Mutated in place (its pruned submodules are
            replaced) and also returned for convenience.
        target_layers: Names (as in `model.named_modules()`) of `Conv2d`/`Linear`
            layers whose output channels/neurons should be pruned.
        sample_input: A representative input tensor, used to trace the model's
            dataflow graph and determine each node's actual tensor shape. Not
            used for any weight-affecting computation.

    Returns:
        nn.Module: The same `model`, with the targeted layers (and any
        dependent BatchNorm / downstream layer) physically shrunk.

    Raises:
        ValueError: If a target layer does not exist, is not a `Conv2d`/`Linear`
            with `groups == 1`, or its output cannot be safely traced to a
            single downstream layer (branching, unsupported op, or an unsafe
            flatten across non-unit spatial dimensions).
    """
    traced = torch.fx.symbolic_trace(model)
    ShapeProp(traced).propagate(sample_input)
    node_by_target = {n.target: n for n in traced.graph.nodes if n.op == "call_module"}
    # Used only to type-check nodes while walking the graph (Conv2d vs BatchNorm
    # vs activation, ...). Structurally identical to `model`'s own module tree,
    # so it stays valid for type-checking even after `model` is mutated below.
    traced_modules = dict(traced.named_modules())

    for layer_name in target_layers:
        self._prune_one_layer(model, layer_name, node_by_target, traced_modules)

    return model

benchmark

benchmark(
    original_model: Module,
    pruned_model: Module,
    sample_input: Tensor,
    original_name: str = "Original (Dense)",
    pruned_name: str = "Pruned (Physically Shrunk)",
    val_dataloader: DataLoader | None = None,
    device: str | device = "cpu",
    compute_flops: bool = False,
) -> BenchmarkReport

Compares the original model against the physically pruned one.

Unlike Pruner.benchmark, no .finalize()-style step is needed first: pruned_model (as returned by .apply()) already has fewer parameters, so real gains in size, latency, and FLOPs are expected here, not just theoretical sparsity.

Parameters:

Name	Type	Description	Default
original_model	Module	The original, unpruned model.	required
pruned_model	Module	The model returned by .apply().	required
sample_input	Tensor	A dummy tensor for latency measurement.	required
original_name	str	Display label for the original model.	'Original (Dense)'
pruned_name	str	Display label for the pruned model.	'Pruned (Physically Shrunk)'
val_dataloader	DataLoader | None	Optional dataloader for accuracy comparison.	None
device	str | device	Device for the benchmark.	'cpu'
compute_flops	bool	If True, also reports FLOPs per sample for both models. Safe to enable here (unlike for Pruner/Quantizer), since physical channel pruning keeps standard Conv2d/Linear ops that FLOPs counting handles accurately. Defaults to False.	False

Returns:

Name	Type	Description
BenchmarkReport	BenchmarkReport	Structured benchmark report ready for .show().

Source code in src/shrinkai/compression/pruning/channel_pruner.py

def benchmark(
    self,
    original_model: nn.Module,
    pruned_model: nn.Module,
    sample_input: torch.Tensor,
    original_name: str = "Original (Dense)",
    pruned_name: str = "Pruned (Physically Shrunk)",
    val_dataloader: DataLoader | None = None,
    device: str | torch.device = "cpu",
    compute_flops: bool = False,
) -> BenchmarkReport:
    """Compares the original model against the physically pruned one.

    Unlike `Pruner.benchmark`, no `.finalize()`-style step is needed first:
    `pruned_model` (as returned by `.apply()`) already has fewer parameters,
    so real gains in size, latency, and FLOPs are expected here, not just
    theoretical sparsity.

    Args:
        original_model: The original, unpruned model.
        pruned_model: The model returned by `.apply()`.
        sample_input: A dummy tensor for latency measurement.
        original_name: Display label for the original model.
        pruned_name: Display label for the pruned model.
        val_dataloader: Optional dataloader for accuracy comparison.
        device: Device for the benchmark.
        compute_flops: If True, also reports FLOPs per sample for both
            models. Safe to enable here (unlike for `Pruner`/`Quantizer`),
            since physical channel pruning keeps standard Conv2d/Linear ops
            that FLOPs counting handles accurately. Defaults to False.

    Returns:
        BenchmarkReport: Structured benchmark report ready for `.show()`.
    """
    original_acc: float | None = None
    pruned_acc: float | None = None

    if val_dataloader is not None:
        original_acc = compute_accuracy(original_model, val_dataloader, device)
        pruned_acc = compute_accuracy(pruned_model, val_dataloader, device)

    return Profiler.compare(
        teacher=original_model,
        student=pruned_model,
        sample_input=sample_input,
        device=device,
        teacher_name=original_name,
        student_name=pruned_name,
        teacher_acc=original_acc,
        student_acc=pruned_acc,
        compute_flops=compute_flops,
    )

Pruner

Orchestrates the pruning of PyTorch models to induce sparsity.

The Pruner traverses the model and applies masking to the weights according to the PruningConfig. It seamlessly supports "Pruning-Aware Training" since PyTorch pruning attaches forward pre-hooks to dynamically mask weights during the forward pass, allowing gradient updates on the remaining unpruned weights.

Equation

Both criteria rank elements/channels by an importance score \(s\) and mask (zero out) the fraction amount with the lowest score:

\[ s = \begin{cases} |w_i| & \text{unstructured (Han et al. (2015)): per-weight magnitude} \\ \|W_j\|_2 = \left(\sum_i w_{j,i}^2\right)^{1/2} & \text{structured: per-output-channel } L_2 \text{ norm} \end{cases} \]

The structured criterion follows the filter-pruning idea of Li et al. (2017), who originally ranked filters by their \(L_1\) norm; this implementation uses PyTorch's torch.nn.utils.prune.ln_structured with \(n=2\) instead. Neither criterion changes tensor shapes (the pruned elements/channels stay in memory as zeros). Refer to ChannelPruner for physical channel removal using the same \(L_2\) criterion.

Parameters:

Name	Type	Description	Default
config	PruningConfig	The configuration object defining the pruning rules.	required

Methods:

Name	Description
apply	Applies the selected pruning strategy to the model's target layers.
benchmark	Evaluates and compares the performance footprint of the dense vs. pruned model.
finalize	Makes the pruning permanent by removing the PyTorch forward hooks and

Source code in src/shrinkai/compression/pruning/pruning.py

class Pruner:
    r"""
    Orchestrates the pruning of PyTorch models to induce sparsity.

    The Pruner traverses the model and applies masking to the weights according
    to the PruningConfig. It seamlessly supports "Pruning-Aware Training" since
    PyTorch pruning attaches forward pre-hooks to dynamically mask weights during
    the forward pass, allowing gradient updates on the remaining unpruned weights.

    Equation:
        Both criteria rank elements/channels by an importance score $s$ and mask
        (zero out) the fraction `amount` with the lowest score:

        $$
        s =
        \begin{cases}
        |w_i| & \text{unstructured (Han et al. (2015)): per-weight magnitude} \\
        \|W_j\|_2 = \left(\sum_i w_{j,i}^2\right)^{1/2} & \text{structured: per-output-channel } L_2 \text{ norm}
        \end{cases}
        $$

        The structured criterion follows the filter-pruning idea of Li et al.
        (2017), who originally ranked filters by their $L_1$ norm; this
        implementation uses PyTorch's `torch.nn.utils.prune.ln_structured` with
        $n=2$ instead. Neither criterion changes tensor shapes (the pruned
        elements/channels stay in memory as zeros). Refer to `ChannelPruner` for
        physical channel removal using the same $L_2$ criterion.

    Args:
        config (PruningConfig): The configuration object defining the pruning rules.
    """  # noqa: E501

    def __init__(self, config: PruningConfig):
        self.config = config

    def apply(self, model: nn.Module) -> nn.Module:
        """
        Applies the selected pruning strategy to the model's target layers.

        Args:
            model (nn.Module): The standard PyTorch model.

        Returns:
            nn.Module: The pruned model (with pruning hooks attached).
        """
        logger.info(f"Applying {self.config.method} pruning ({self.config.amount * 100}%)...")

        for _, module in model.named_modules():
            if isinstance(module, self.config.target_types):
                if self.config.method == "unstructured":
                    self._apply_unstructured(module)
                elif self.config.method == "structured":
                    self._apply_structured(module)

        return model

    def _apply_unstructured(self, module: nn.Module):
        """Internal method: Applies L1 Unstructured Pruning to a specific module."""
        prune.l1_unstructured(module, name="weight", amount=self.config.amount)
        if getattr(module, "bias", None) is not None:
            prune.l1_unstructured(module, name="bias", amount=self.config.amount)

    def _apply_structured(self, module: nn.Module):
        """Internal method: Applies L2 Structured Pruning (e.g., removing output channels)."""
        prune.ln_structured(module, name="weight", amount=self.config.amount, n=2, dim=0)

    @staticmethod
    def finalize(pruned_model: nn.Module) -> nn.Module:
        """
        Makes the pruning permanent by removing the PyTorch forward hooks and
        baking the sparsity mask directly into the weight tensors.

        This MUST be called after training/distillation before saving the model
        for deployment, otherwise the masking overhead slows down inference.

        Args:
            pruned_model (nn.Module): The model with pruning hooks attached.

        Returns:
            nn.Module: The permanent, clean sparse model.
        """
        if getattr(pruned_model, "_pruning_finalized", False):
            logger.warning("Pruning has already been finalized on this model. Skipping.")
            return pruned_model

        logger.info("Finalizing pruning: removing hooks and baking masks into weights...")
        hooks_removed = False
        for _, module in pruned_model.named_modules():
            if hasattr(module, "weight_mask"):
                prune.remove(module, "weight")
                hooks_removed = True
            if hasattr(module, "bias_mask"):
                prune.remove(module, "bias")
                hooks_removed = True

        if not hooks_removed:
            logger.warning("No pruning hooks were found. Was the model actually pruned?")
        pruned_model._pruning_finalized = True

        return pruned_model

    def benchmark(
        self,
        original_model: nn.Module,
        pruned_model: nn.Module,
        sample_input: torch.Tensor,
        original_name: str = "Original (Dense)",
        pruned_name: str = "Pruned (Sparse)",
        val_dataloader: DataLoader | None = None,
        device: str | torch.device = "cpu",
    ) -> BenchmarkReport:
        """
        Evaluates and compares the performance footprint of the dense vs. pruned model.
        Note: True latency gains for unstructured pruning require sparse-tensor
        supported hardware engines.

        Args:
            original_model (nn.Module): The dense baseline model.
            pruned_model (nn.Module): The pruned model.
            sample_input (torch.Tensor): A dummy tensor for latency measurement.
            original_name (str, optional): Display label for baseline.
            pruned_name (str, optional): Display label for pruned model.
            val_dataloader (DataLoader | None, optional): Dataloader for accuracy.
            device (str | torch.device, optional): Device for the benchmark.

        Returns:
                    BenchmarkReport: Structured benchmark report ready for `.show()`.
        """
        eval_pruned_model = _clone_module(pruned_model)
        eval_pruned_model = self.finalize(eval_pruned_model)

        original_acc: float | None = None
        pruned_acc: float | None = None

        if val_dataloader is not None:
            original_acc = compute_accuracy(original_model, val_dataloader, device)
            pruned_acc = compute_accuracy(eval_pruned_model, val_dataloader, device)

        return Profiler.compare(
            teacher=original_model,
            student=eval_pruned_model,
            sample_input=sample_input,
            device=device,
            teacher_name=original_name,
            student_name=pruned_name,
            teacher_acc=original_acc,
            student_acc=pruned_acc,
        )

Methods:

apply

apply(model: Module) -> nn.Module

Applies the selected pruning strategy to the model's target layers.

Parameters:

Name	Type	Description	Default
model	Module	The standard PyTorch model.	required

Returns:

Type	Description
Module	nn.Module: The pruned model (with pruning hooks attached).

Source code in src/shrinkai/compression/pruning/pruning.py

def apply(self, model: nn.Module) -> nn.Module:
    """
    Applies the selected pruning strategy to the model's target layers.

    Args:
        model (nn.Module): The standard PyTorch model.

    Returns:
        nn.Module: The pruned model (with pruning hooks attached).
    """
    logger.info(f"Applying {self.config.method} pruning ({self.config.amount * 100}%)...")

    for _, module in model.named_modules():
        if isinstance(module, self.config.target_types):
            if self.config.method == "unstructured":
                self._apply_unstructured(module)
            elif self.config.method == "structured":
                self._apply_structured(module)

    return model

benchmark

benchmark(
    original_model: Module,
    pruned_model: Module,
    sample_input: Tensor,
    original_name: str = "Original (Dense)",
    pruned_name: str = "Pruned (Sparse)",
    val_dataloader: DataLoader | None = None,
    device: str | device = "cpu",
) -> BenchmarkReport

Evaluates and compares the performance footprint of the dense vs. pruned model. Note: True latency gains for unstructured pruning require sparse-tensor supported hardware engines.

Parameters:

Name	Type	Description	Default
original_model	Module	The dense baseline model.	required
pruned_model	Module	The pruned model.	required
sample_input	Tensor	A dummy tensor for latency measurement.	required
original_name	str	Display label for baseline.	'Original (Dense)'
pruned_name	str	Display label for pruned model.	'Pruned (Sparse)'
val_dataloader	DataLoader | None	Dataloader for accuracy.	None
device	str | device	Device for the benchmark.	'cpu'

Returns:

Name	Type	Description
BenchmarkReport	BenchmarkReport	Structured benchmark report ready for .show().

Source code in src/shrinkai/compression/pruning/pruning.py

def benchmark(
    self,
    original_model: nn.Module,
    pruned_model: nn.Module,
    sample_input: torch.Tensor,
    original_name: str = "Original (Dense)",
    pruned_name: str = "Pruned (Sparse)",
    val_dataloader: DataLoader | None = None,
    device: str | torch.device = "cpu",
) -> BenchmarkReport:
    """
    Evaluates and compares the performance footprint of the dense vs. pruned model.
    Note: True latency gains for unstructured pruning require sparse-tensor
    supported hardware engines.

    Args:
        original_model (nn.Module): The dense baseline model.
        pruned_model (nn.Module): The pruned model.
        sample_input (torch.Tensor): A dummy tensor for latency measurement.
        original_name (str, optional): Display label for baseline.
        pruned_name (str, optional): Display label for pruned model.
        val_dataloader (DataLoader | None, optional): Dataloader for accuracy.
        device (str | torch.device, optional): Device for the benchmark.

    Returns:
                BenchmarkReport: Structured benchmark report ready for `.show()`.
    """
    eval_pruned_model = _clone_module(pruned_model)
    eval_pruned_model = self.finalize(eval_pruned_model)

    original_acc: float | None = None
    pruned_acc: float | None = None

    if val_dataloader is not None:
        original_acc = compute_accuracy(original_model, val_dataloader, device)
        pruned_acc = compute_accuracy(eval_pruned_model, val_dataloader, device)

    return Profiler.compare(
        teacher=original_model,
        student=eval_pruned_model,
        sample_input=sample_input,
        device=device,
        teacher_name=original_name,
        student_name=pruned_name,
        teacher_acc=original_acc,
        student_acc=pruned_acc,
    )

finalize staticmethod

finalize(pruned_model: Module) -> nn.Module

Makes the pruning permanent by removing the PyTorch forward hooks and baking the sparsity mask directly into the weight tensors.

This MUST be called after training/distillation before saving the model for deployment, otherwise the masking overhead slows down inference.

Parameters:

Name	Type	Description	Default
pruned_model	Module	The model with pruning hooks attached.	required

Returns:

Type	Description
Module	nn.Module: The permanent, clean sparse model.

Source code in src/shrinkai/compression/pruning/pruning.py

@staticmethod
def finalize(pruned_model: nn.Module) -> nn.Module:
    """
    Makes the pruning permanent by removing the PyTorch forward hooks and
    baking the sparsity mask directly into the weight tensors.

    This MUST be called after training/distillation before saving the model
    for deployment, otherwise the masking overhead slows down inference.

    Args:
        pruned_model (nn.Module): The model with pruning hooks attached.

    Returns:
        nn.Module: The permanent, clean sparse model.
    """
    if getattr(pruned_model, "_pruning_finalized", False):
        logger.warning("Pruning has already been finalized on this model. Skipping.")
        return pruned_model

    logger.info("Finalizing pruning: removing hooks and baking masks into weights...")
    hooks_removed = False
    for _, module in pruned_model.named_modules():
        if hasattr(module, "weight_mask"):
            prune.remove(module, "weight")
            hooks_removed = True
        if hasattr(module, "bias_mask"):
            prune.remove(module, "bias")
            hooks_removed = True

    if not hooks_removed:
        logger.warning("No pruning hooks were found. Was the model actually pruned?")
    pruned_model._pruning_finalized = True

    return pruned_model

PruningConfig dataclass

Configuration parameters for model pruning.

Attributes:

Name	Type	Description
method	str	The pruning strategy to apply. - "unstructured": Removes individual weights based on L1-norm (closest to zero). Creates sparse tensors but doesn't change tensor shapes. - "structured": Zeroes entire channels/neurons based on their Ln-norm, via torch.nn.utils.prune. This does NOT change tensor shapes either: the pruned channels stay in memory as zeros, so it does not by itself reduce parameter count, model size, or inference latency on standard hardware. To physically remove channels (actually shrinking the model), use ChannelPruner instead. Defaults to "unstructured".
amount	float	The fraction of weights/channels to prune. Must be a float between 0.0 and 1.0 (e.g., 0.3 means 30% pruned). Defaults to 0.3.
target_types	tuple	The PyTorch module types to apply pruning to. Defaults to (nn.Linear, nn.Conv2d).

Source code in src/shrinkai/compression/pruning/pruning.py

@dataclass
class PruningConfig:
    """
    Configuration parameters for model pruning.

    Attributes:
        method (str): The pruning strategy to apply.
            - "unstructured": Removes individual weights based on L1-norm (closest to zero).
              Creates sparse tensors but doesn't change tensor shapes.
            - "structured": Zeroes entire channels/neurons based on their Ln-norm,
              via `torch.nn.utils.prune`. This does NOT change tensor shapes either:
              the pruned channels stay in memory as zeros, so it does not by itself
              reduce parameter count, model size, or inference latency on standard
              hardware. To physically remove channels (actually shrinking the
              model), use `ChannelPruner` instead.
            Defaults to "unstructured".

        amount (float): The fraction of weights/channels to prune.
            Must be a float between 0.0 and 1.0 (e.g., 0.3 means 30% pruned).
            Defaults to 0.3.

        target_types (tuple): The PyTorch module types to apply pruning to.
            Defaults to (nn.Linear, nn.Conv2d).
    """

    method: str = "unstructured"
    amount: float = 0.3
    target_types: tuple = (nn.Linear, nn.Conv2d)

    def __post_init__(self):
        valid_methods = ["unstructured", "structured"]
        if self.method not in valid_methods:
            raise ValueError(f"Invalid method '{self.method}'. Supported choices: {valid_methods}")
        if not (0.0 < self.amount < 1.0):
            raise ValueError(f"Pruning amount must be between 0.0 and 1.0, got {self.amount}")