compression

shrinkai.compression

Model compression: pruning and quantization.

Two independent, composable families of techniques for shrinking a model after (or during) training:

• Pruning (shrinkai.compression.pruning): induces sparsity (Pruner) or physically removes channels (ChannelPruner).
• Quantization (shrinkai.compression.quantization): reduces numerical precision of weights/activations (Quantizer), via PTQ or QAT.

Both can be combined with shrinkai.distillation (e.g. quantization-aware training driven by a Distiller) or used standalone on any nn.Module.

Modules:

Name	Description
pruning	Pruning: inducing sparsity or physically shrinking a model.
quantization	Quantization: reducing the numerical precision of a model's weights/activations.

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.
QuantConfig	Configuration parameters for model quantization.
Quantizer	Orchestrates the quantization of PyTorch models.

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}")

QuantConfig dataclass

Configuration parameters for model quantization.

This dataclass standardizes how quantization is applied across different backends and strategies. It defines the target bit-width and the mathematical approach used to compress the network.

Attributes:

Name	Type	Description
target_dtype	str	The target data type for the model's weights. Currently supported options include: - "int8": 8-bit integer (Standard for CPU/Edge deployment). - "fp16": 16-bit float (Standard for GPU memory reduction). Defaults to "int8".
strategy	str	The quantization strategy to apply. - "ptq" (Post-Training Quantization): Applies immediate mathematical conversion to the weights. No gradient computation or training is required. - "qat" (Quantization-Aware Training): Prepares the model with FakeQuantize nodes. Requires subsequent training (e.g., via shrinkai.Distiller) before final conversion. Defaults to "ptq".
backend	str	The underlying engine executing the quantization. - "torch": Native PyTorch quantization (fbgemm/qnnpack). Ideal for CNNs and small LMs. - "bitsandbytes": (Reserved for future LLM integration) Block-wise quantization. Defaults to "torch".
calibrate_data	Any | None	A dataloader (or any iterable of batches, each either a plain input tensor or a (inputs, labels, ...) tuple/list) used exclusively for Static PTQ. If provided, the quantizer runs a forward pass over this data to calibrate activation scales before conversion. If None, Dynamic PTQ is used instead (weights only, no calibration). Defaults to None.

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

@dataclass
class QuantConfig:
    """
    Configuration parameters for model quantization.

    This dataclass standardizes how quantization is applied across different
    backends and strategies. It defines the target bit-width and the mathematical
    approach used to compress the network.

    Attributes:
        target_dtype (str): The target data type for the model's weights.
            Currently supported options include:
            - "int8": 8-bit integer (Standard for CPU/Edge deployment).
            - "fp16": 16-bit float (Standard for GPU memory reduction).
            Defaults to "int8".

        strategy (str): The quantization strategy to apply.
            - "ptq" (Post-Training Quantization): Applies immediate mathematical
              conversion to the weights. No gradient computation or training is required.
            - "qat" (Quantization-Aware Training): Prepares the model with FakeQuantize
              nodes. Requires subsequent training (e.g., via `shrinkai.Distiller`)
              before final conversion.
            Defaults to "ptq".

        backend (str): The underlying engine executing the quantization.
            - "torch": Native PyTorch quantization (fbgemm/qnnpack). Ideal for CNNs and small LMs.
            - "bitsandbytes": (Reserved for future LLM integration) Block-wise quantization.
            Defaults to "torch".

        calibrate_data (Any | None): A dataloader (or any iterable of batches, each
            either a plain input tensor or a `(inputs, labels, ...)` tuple/list) used
            exclusively for Static PTQ. If provided, the quantizer runs a forward
            pass over this data to calibrate activation scales before conversion.
            If None, Dynamic PTQ is used instead (weights only, no calibration).
            Defaults to None.
    """

    target_dtype: str = "int8"
    strategy: str = "ptq"
    backend: str = "torch"
    calibrate_data: Any | None = None

    def __post_init__(self):
        valid_strategies = ["ptq", "qat"]
        valid_dtype = ["int8", "fp16"]

        if self.strategy not in valid_strategies:
            raise ValueError(
                f"Invalid strategy '{self.strategy}'. Supported choices: {valid_strategies}"
            )

        if self.target_dtype not in valid_dtype:
            raise ValueError(
                f"Invalid strategy '{self.target_dtype}'. Supported choices: {valid_dtype}"
            )

Quantizer

Orchestrates the quantization of PyTorch models.

The Quantizer reads a QuantConfig and safely modifies the computational graph of a given PyTorch nn.Module. It handles the complexities of PyTorch's native quantization APIs.

Equation

PTQ and QAT both rely on the affine (uniform) quantization scheme of Jacob et al. (2018):

\[q = \text{clip}\left(\text{round}\left(\frac{x}{s}\right) + z,\ q_{min},\ q_{max}\right), \qquad x \approx s \, (q - z)\]

where \(s > 0\) (scale) and \(z\) (zero-point) are derived from the observed range of \(x\), via calibration on calibrate_data for static PTQ activations, on the fly per-batch for dynamic PTQ, or from weight statistics directly, and \([q_{min}, q_{max}]\) bounds the target integer range (e.g. \([-128, 127]\) for int8). PyTorch's default backends use \(z = 0\) (symmetric) for weights and the full affine form for activations.

For QAT, this same round-trip is simulated in the forward pass ("Fake Quantization") while gradients flow through it via the straight-through estimator (\(\partial q / \partial x \approx 1\) inside \([x_{min}, x_{max}]\), \(0\) outside), letting the model adapt its weights to the quantization noise before finalize_qat() converts it to real integer weights.

Parameters:

Name	Type	Description	Default
config	QuantConfig	The configuration object defining the quantization rules.	required

Methods:

Name	Description
apply	Applies the selected quantization strategy to the model.
benchmark	Runs complete profiling suite on original and quantized models.
finalize_qat	To be called AFTER the distillation training loop (distiller.fit).

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

class Quantizer:
    r"""
    Orchestrates the quantization of PyTorch models.

    The Quantizer reads a `QuantConfig` and safely modifies the computational graph
    of a given PyTorch `nn.Module`. It handles the complexities of PyTorch's native
    quantization APIs.

    Equation:
        PTQ and QAT both rely on the affine (uniform) quantization scheme of
        Jacob et al. (2018):

        $$q = \text{clip}\left(\text{round}\left(\frac{x}{s}\right) + z,\ q_{min},\ q_{max}\right), \qquad x \approx s \, (q - z)$$

        where $s > 0$ (scale) and $z$ (zero-point) are derived from the observed
        range of $x$, via calibration on `calibrate_data` for static PTQ
        activations, on the fly per-batch for dynamic PTQ, or from weight
        statistics directly, and $[q_{min}, q_{max}]$ bounds the target integer
        range (e.g. $[-128, 127]$ for int8). PyTorch's default backends use
        $z = 0$ (symmetric) for weights and the full affine form for activations.

        For QAT, this same round-trip is simulated in the forward pass ("Fake
        Quantization") while gradients flow through it via the straight-through
        estimator ($\partial q / \partial x \approx 1$ inside $[x_{min}, x_{max}]$,
        $0$ outside), letting the model adapt its weights to the quantization
        noise before `finalize_qat()` converts it to real integer weights.

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

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

    def apply(self, model: nn.Module) -> nn.Module:
        """
        Applies the selected quantization strategy to the model.

        Depending on `config.strategy`, this method will either immediately convert
        the weights (PTQ) or insert FakeQuantize nodes for future training (QAT).

        Args:
            model (nn.Module): The standard, full-precision PyTorch model.

        Returns:
            nn.Module: The modified model. If PTQ, it is ready for deployment.
            If QAT, it must be trained and then passed to `finalize_qat()`.
        """
        if self.config.strategy == "ptq":
            return self._apply_ptq(model)
        elif self.config.strategy == "qat":
            return self._apply_qat(model)
        else:
            raise NotImplementedError(f"Strategy {self.config.strategy} is not implemented yet.")

    def _apply_ptq(self, model: nn.Module) -> nn.Module:
        """
        Applies Post-Training Quantization (PTQ).

        Uses Static PTQ (calibrated on `config.calibrate_data`) if provided,
        otherwise falls back to Dynamic PTQ (weights only, no calibration needed).
        Ideal for immediate inference optimization without retraining.
        """
        logger.info(
            f"Applying PTQ ({self.config.target_dtype}) using {self.config.backend} backend..."
        )

        if self.config.backend != "torch" or self.config.target_dtype != "int8":
            raise NotImplementedError(
                f"PTQ is not yet supported for backend '{self.config.backend}' "
                f"with dtype '{self.config.target_dtype}'."
            )

        if self.config.calibrate_data is not None:
            return self._apply_static_ptq(model)

        return torch.ao.quantization.quantize_dynamic(model, {nn.Linear}, dtype=torch.qint8)

    def _apply_static_ptq(self, model: nn.Module) -> nn.Module:
        """
        Applies Static Post-Training Quantization, calibrated on `config.calibrate_data`.

        Wraps the model with `QuantStub`/`DeQuantStub` (via `QuantWrapper`) so it keeps
        accepting and returning standard float tensors, runs a calibration pass over
        `config.calibrate_data` to observe activation ranges, then converts both
        weights and activations to int8.

        Note:
            For best accuracy, fuse Conv-BN-ReLU sequences on `model` (via
            `torch.ao.quantization.fuse_modules`) before calling `apply()`. This
            implementation works without fusion, at the cost of some accuracy.
        """
        logger.info("Calibrating activations for Static PTQ...")

        wrapped_model = torch.ao.quantization.QuantWrapper(model)
        wrapped_model.eval()

        engine = torch.backends.quantized.engine
        wrapped_model.qconfig = torch.ao.quantization.get_default_qconfig(engine)
        prepared_model = torch.ao.quantization.prepare(wrapped_model, inplace=False)

        with torch.no_grad():
            for batch in self.config.calibrate_data:
                inputs = batch[0] if isinstance(batch, list | tuple) else batch
                prepared_model(inputs)

        return torch.ao.quantization.convert(prepared_model, inplace=False)

    def _apply_qat(self, model: nn.Module) -> nn.Module:
        """
        Prepares the model for Quantization-Aware Training (QAT).
        Inserts 'Fake Quantization' nodes without converting the actual tensors,
        allowing gradients to flow during the distillation process.
        """
        logger.info(f"Preparing model for QAT ({self.config.target_dtype})...")

        if self.config.backend == "torch":
            model.train()
            engine = torch.backends.quantized.engine
            model.qconfig = torch.ao.quantization.get_default_qat_qconfig(engine)
            qat_model = torch.ao.quantization.prepare_qat(model, inplace=False)
            return qat_model
        else:
            raise NotImplementedError(
                f"QAT is not yet supported for backend '{self.config.backend}'."
            )

    @staticmethod
    def finalize_qat(qat_model: nn.Module) -> nn.Module:
        """
        To be called AFTER the distillation training loop (distiller.fit).
        Converts the simulated FakeQuantize nodes into actual quantized weights (e.g., int8).

        Args:
            qat_model (nn.Module): The model trained with FakeQuantize nodes.

        Returns:
            nn.Module: The fully quantized model.
        """
        logger.info(
            "Finalizing QAT model: converting FakeQuantize nodes to actual quantized weights..."
        )
        qat_model.eval()
        return torch.ao.quantization.convert(qat_model, inplace=False)

    def benchmark(
        self,
        original_model: nn.Module,
        quantized_model: nn.Module,
        sample_input: torch.Tensor,
        original_name: str = "Original (FP32)",
        quantized_name: str = "Quantized",
        val_dataloader: DataLoader | None = None,
        device: str | torch.device = "cpu",
    ) -> BenchmarkReport:
        """Runs complete profiling suite on original and quantized models.

        Args:
            original_model: The base FP32 PyTorch model.
            quantized_model: The model returned by `.apply()`.
            sample_input: Batch tensor matching target inference dimension.
            original_name: Display label for the original model.
            quantized_name: Display label for the quantized model.
            val_dataloader: Optional dataloader to compute final accuracy metrics.
            device: Device to run the benchmark on (default "cpu", as INT8 is often CPU-optimized).

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

        eval_quantized_model = quantized_model
        if self.config.strategy == "qat":
            eval_quantized_model = self.finalize_qat(quantized_model)

        if val_dataloader is not None:
            original_acc = compute_accuracy(original_model, val_dataloader, device)
            quantized_acc = compute_accuracy(eval_quantized_model, val_dataloader, device)

        return Profiler.compare(
            teacher=original_model,
            student=eval_quantized_model,
            sample_input=sample_input,
            device=device,
            teacher_name=original_name,
            student_name=quantized_name,
            teacher_acc=original_acc,
            student_acc=quantized_acc,
        )

Methods:

apply

apply(model: Module) -> nn.Module

Applies the selected quantization strategy to the model.

Depending on config.strategy, this method will either immediately convert the weights (PTQ) or insert FakeQuantize nodes for future training (QAT).

Parameters:

Name	Type	Description	Default
model	Module	The standard, full-precision PyTorch model.	required

Returns:

Type	Description
Module	nn.Module: The modified model. If PTQ, it is ready for deployment.
Module	If QAT, it must be trained and then passed to finalize_qat().

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

def apply(self, model: nn.Module) -> nn.Module:
    """
    Applies the selected quantization strategy to the model.

    Depending on `config.strategy`, this method will either immediately convert
    the weights (PTQ) or insert FakeQuantize nodes for future training (QAT).

    Args:
        model (nn.Module): The standard, full-precision PyTorch model.

    Returns:
        nn.Module: The modified model. If PTQ, it is ready for deployment.
        If QAT, it must be trained and then passed to `finalize_qat()`.
    """
    if self.config.strategy == "ptq":
        return self._apply_ptq(model)
    elif self.config.strategy == "qat":
        return self._apply_qat(model)
    else:
        raise NotImplementedError(f"Strategy {self.config.strategy} is not implemented yet.")

benchmark

benchmark(
    original_model: Module,
    quantized_model: Module,
    sample_input: Tensor,
    original_name: str = "Original (FP32)",
    quantized_name: str = "Quantized",
    val_dataloader: DataLoader | None = None,
    device: str | device = "cpu",
) -> BenchmarkReport

Runs complete profiling suite on original and quantized models.

Parameters:

Name	Type	Description	Default
original_model	Module	The base FP32 PyTorch model.	required
quantized_model	Module	The model returned by .apply().	required
sample_input	Tensor	Batch tensor matching target inference dimension.	required
original_name	str	Display label for the original model.	'Original (FP32)'
quantized_name	str	Display label for the quantized model.	'Quantized'
val_dataloader	DataLoader | None	Optional dataloader to compute final accuracy metrics.	None
device	str | device	Device to run the benchmark on (default "cpu", as INT8 is often CPU-optimized).	'cpu'

Returns:

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

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

def benchmark(
    self,
    original_model: nn.Module,
    quantized_model: nn.Module,
    sample_input: torch.Tensor,
    original_name: str = "Original (FP32)",
    quantized_name: str = "Quantized",
    val_dataloader: DataLoader | None = None,
    device: str | torch.device = "cpu",
) -> BenchmarkReport:
    """Runs complete profiling suite on original and quantized models.

    Args:
        original_model: The base FP32 PyTorch model.
        quantized_model: The model returned by `.apply()`.
        sample_input: Batch tensor matching target inference dimension.
        original_name: Display label for the original model.
        quantized_name: Display label for the quantized model.
        val_dataloader: Optional dataloader to compute final accuracy metrics.
        device: Device to run the benchmark on (default "cpu", as INT8 is often CPU-optimized).

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

    eval_quantized_model = quantized_model
    if self.config.strategy == "qat":
        eval_quantized_model = self.finalize_qat(quantized_model)

    if val_dataloader is not None:
        original_acc = compute_accuracy(original_model, val_dataloader, device)
        quantized_acc = compute_accuracy(eval_quantized_model, val_dataloader, device)

    return Profiler.compare(
        teacher=original_model,
        student=eval_quantized_model,
        sample_input=sample_input,
        device=device,
        teacher_name=original_name,
        student_name=quantized_name,
        teacher_acc=original_acc,
        student_acc=quantized_acc,
    )

finalize_qat staticmethod

finalize_qat(qat_model: Module) -> nn.Module

To be called AFTER the distillation training loop (distiller.fit). Converts the simulated FakeQuantize nodes into actual quantized weights (e.g., int8).

Parameters:

Name	Type	Description	Default
qat_model	Module	The model trained with FakeQuantize nodes.	required

Returns:

Type	Description
Module	nn.Module: The fully quantized model.

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

@staticmethod
def finalize_qat(qat_model: nn.Module) -> nn.Module:
    """
    To be called AFTER the distillation training loop (distiller.fit).
    Converts the simulated FakeQuantize nodes into actual quantized weights (e.g., int8).

    Args:
        qat_model (nn.Module): The model trained with FakeQuantize nodes.

    Returns:
        nn.Module: The fully quantized model.
    """
    logger.info(
        "Finalizing QAT model: converting FakeQuantize nodes to actual quantized weights..."
    )
    qat_model.eval()
    return torch.ao.quantization.convert(qat_model, inplace=False)