channel_pruner

shrinkai.compression.pruning.channel_pruner

Physical (dependency-aware) structured channel pruning.

Unlike Pruner (which only zeroes out weights via torch.nn.utils.prune masks, without changing tensor shapes or producing real latency/FLOPs gains), ChannelPruner actually rebuilds smaller Conv2d/Linear (and dependent BatchNorm) layers with the pruned channels physically removed.

Determining which layers can be safely shrunk together requires knowing the model's real dataflow graph, removing a layer's output channels is only valid if every downstream consumer of that output has its input channels shrunk to match. This module uses torch.fx to trace that graph and torch.fx.passes.shape_prop.ShapeProp to know each node's actual tensor shape.

Scope (by design, to stay correct rather than merely "not crashing"): - Supports simple, non-branching chains: a targeted Conv2d/Linear layer whose output feeds, without branching (exactly one consumer at every hop), through any number of BatchNorm/activation/pooling/dropout layers, into either another Conv2d/Linear layer or the model's final output. - A flatten/view/reshape between a 4D conv-style tensor and a Linear layer is only allowed once the spatial dimensions have already been reduced to 1x1 (e.g. by AdaptiveAvgPool2d(1)). At that point each channel maps to exactly one flattened feature, so no interleaving ambiguity exists. - Grouped/depthwise convolutions (groups != 1) are rejected: removing channels from a grouped conv can change which input channels feed which output channels, which this implementation does not attempt to resolve. - Any branching topology (residual/skip connections, concatenation, attention, multiple consumers of the same tensor) is rejected with a clear error rather than silently producing an incorrect model. For those architectures, use Pruner (mask-based, safe on any topology, but without physical shrinkage) instead.

Classes:

Name	Description
`ChannelPruner`	Physically removes pruned output channels from `Conv2d`/`Linear` layers.

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

channel_pruner

shrinkai.compression.pruning.channel_pruner

Classes

ChannelPruner

Methods:

apply

benchmark

Functions: