From fcfa686959dcaa88ad93cf0696478ce323cfb274 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ey=C3=BCp=20Can=20Akman?= Date: Sat, 18 Jul 2026 14:17:13 +0300 Subject: [PATCH] fix(pruning): normalize a negative channel axis ChannelStructured(axis=-1) pruned the wrong channels. _compute_channel_mask compared each dim index against the raw axis when building its reduce list, so a negative axis excluded nothing and the per-channel L1 norms collapsed to a scalar. Normalize the axis first, matching PerChannelGranularity, and reject an out-of-range axis with ValueError rather than letting it reach the reduction. --- src/coreai_opt/pruning/spec/prune.py | 7 +++ src/coreai_opt/pruning/spec/scheme.py | 3 ++ tests/pruning/test_magnitude_pruner.py | 71 +++++++++++++++++++++++++- 3 files changed, 79 insertions(+), 2 deletions(-) diff --git a/src/coreai_opt/pruning/spec/prune.py b/src/coreai_opt/pruning/spec/prune.py index 12c1e77..5ddf2bf 100644 --- a/src/coreai_opt/pruning/spec/prune.py +++ b/src/coreai_opt/pruning/spec/prune.py @@ -170,6 +170,13 @@ def _compute_channel_mask( Channel importance is measured by L1 norm. The least-important channels are pruned entirely. """ + if not (-weight.ndim <= axis < weight.ndim): + raise ValueError( + f"Invalid axis. Should be in range [{-weight.ndim}, {weight.ndim}), but got {axis}" + ) + if axis < 0: + axis += weight.ndim + num_channels = weight.shape[axis] num_prune = math.floor(num_channels * sparsity) diff --git a/src/coreai_opt/pruning/spec/scheme.py b/src/coreai_opt/pruning/spec/scheme.py index 7697619..3049db6 100644 --- a/src/coreai_opt/pruning/spec/scheme.py +++ b/src/coreai_opt/pruning/spec/scheme.py @@ -70,6 +70,9 @@ class ChannelStructured(PruningScheme): Entire channels (slices along ``axis``) are pruned or kept together. Channel importance is determined by the pruning algorithm (e.g. L1 norm of each channel for magnitude-based pruning). + + Note: + ``axis`` can be negatively indexed as per standard Python style indexing. """ axis: int = Field(default=0, description="Axis along which channels are pruned.") diff --git a/tests/pruning/test_magnitude_pruner.py b/tests/pruning/test_magnitude_pruner.py index 35aef7a..b3b60e3 100644 --- a/tests/pruning/test_magnitude_pruner.py +++ b/tests/pruning/test_magnitude_pruner.py @@ -369,7 +369,8 @@ def test_channel_structured_pruning_hand(self) -> None: ) assert torch.equal(model.weight.detach(), expected) - def test_channel_structured_conv2d(self) -> None: + @pytest.mark.parametrize("axis", [0, -4]) + def test_channel_structured_conv2d(self, axis: int) -> None: """Channel-structured pruning on Conv2d zeros entire output filters.""" torch.manual_seed(42) model = nn.Conv2d(3, 8, kernel_size=3, bias=False) @@ -379,7 +380,7 @@ def test_channel_structured_conv2d(self) -> None: op_state_spec={ "weight": PruningSpec( target_sparsity=0.5, - pruning_scheme=ChannelStructured(axis=0), + pruning_scheme=ChannelStructured(axis=axis), ) } ) @@ -395,6 +396,72 @@ def test_channel_structured_conv2d(self) -> None: filt = weight[i] assert filt.eq(0).all() or filt.ne(0).all(), f"Filter {i} is partially pruned" + @pytest.mark.parametrize("target_sparsity", [0.5, 0.75]) + @pytest.mark.parametrize( + "negative_axis,positive_axis", + [(-1, 1), (-2, 0)], + ids=["last-dim", "first-dim"], + ) + def test_channel_structured_negative_axis( + self, negative_axis: int, positive_axis: int, target_sparsity: float + ) -> None: + """A negative axis prunes the same channels as its positive equivalent.""" + + def prune_along(axis: int) -> torch.Tensor: + model = nn.Linear(4, 4, bias=False) + with torch.no_grad(): + model.weight.copy_( + torch.tensor( + [ + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + [9.0, 10.0, 11.0, 12.0], + [13.0, 14.0, 15.0, 16.0], + ] + ) + ) + + config = MagnitudePrunerConfig( + global_config=ModuleMagnitudePrunerConfig( + op_state_spec={ + "weight": PruningSpec( + target_sparsity=target_sparsity, + pruning_scheme=ChannelStructured(axis=axis), + ) + } + ) + ) + pruner = MagnitudePruner(model, config) + pruner.prepare((torch.randn(1, 4),)) + return model.weight.detach() + + pruned = prune_along(negative_axis) + assert torch.equal(pruned, prune_along(positive_axis)) + + # L1 norms increase with index along both axes, so the highest indices survive. + num_keep = 4 - int(4 * target_sparsity) + kept = [i for i in range(4) if pruned.select(positive_axis, i).ne(0).any()] + assert kept == list(range(4 - num_keep, 4)) + + @pytest.mark.parametrize("axis", [-3, 2], ids=["below-range", "above-range"]) + def test_channel_structured_axis_out_of_range(self, axis: int) -> None: + """An axis outside [-ndim, ndim) raises ValueError.""" + model = nn.Linear(4, 4, bias=False) + config = MagnitudePrunerConfig( + global_config=ModuleMagnitudePrunerConfig( + op_state_spec={ + "weight": PruningSpec( + target_sparsity=0.5, + pruning_scheme=ChannelStructured(axis=axis), + ) + } + ) + ) + pruner = MagnitudePruner(model, config) + + with pytest.raises(ValueError, match="Invalid axis"): + pruner.prepare((torch.randn(1, 4),)) + def test_linear_unstructured_conv2d_channel_structured(self) -> None: """Apply unstructured to Linear and channel-structured to Conv2d in same model."""