Skip to content

fix(regularize): norm=2 group lasso in RegularizeCallback, honest docs #97

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nathanhubens merged 1 commit into from
Sep 24, 2026
Merged

fix(regularize): norm=2 group lasso in RegularizeCallback, honest docs #97

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

### New Features
- `Pruner(..., reg=1e-4)` builds torch-pruning's `GroupNormPruner` instead of its `MetaPruner`, and the new `Pruner.regularize(scale=1.)` adds DepGraph's group penalty to the gradients already computed: every channel of a dependency group receives `reg * gamma * w` on each tensor it spans — the convolution's filter, the batch-norm weight and the input slices of the layers that read it. `gamma` runs from `1` on the group's most important channel to `2**alpha` on its least important one (`alpha=4` by default), measured on the square root of the pruner's own `criteria` score. It is called after `backward` and before the optimizer step, and `scale` multiplies the penalty for gradients that are themselves scaled (a mixed-precision loss scale, gradient accumulation). Only the groups the pruner will prune are penalized, so a per-layer dict leaves the layers it does not target alone; a group with a constant importance (a single channel among them) is skipped instead of producing `nan` gradients, with one warning if every group is; the groups are rebuilt after each `prune_model()`. Biases, frozen parameters and a depthwise convolution's own kernels are not penalized. `reg` and `alpha` are keyword-only and `reg < 0` is refused. `reg=0`, the default, builds the same `MetaPruner` as before and prunes identically apart from the device fix below. `PruneCallback` is unchanged
- `RegularizeCallback(norm=2)` penalizes the L2 norm of each group of the chosen `granularity` (group lasso) instead of every weight separately. The default, `norm=1`, is unchanged byte for byte and is an elementwise L1 whatever the granularity, since the groups partition the weights; groups are taken within one layer and do not follow cross-layer couplings such as the next layer's input channels or BatchNorm parameters

### Bug Fixes
- `Pruner.group_importance` returns its scores on the device of the layers it scored instead of fastai's default device, which put them on the GPU for a model on the CPU whenever one was available. The ranking a prune step makes is unchanged: torch-pruning moves the scores to the CPU before ranking them
Expand Down
2 changes: 2 additions & 0 deletions fasterai/_modidx.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -644,6 +644,8 @@
'fasterai/regularize/regularize_callback.py'),
'fasterai.regularize.regularize_callback.RegularizeCallback.__init__': ( 'regularize/regularize_callback.html#regularizecallback.__init__',
'fasterai/regularize/regularize_callback.py'),
'fasterai.regularize.regularize_callback.RegularizeCallback._group_norms': ( 'regularize/regularize_callback.html#regularizecallback._group_norms',
'fasterai/regularize/regularize_callback.py'),
'fasterai.regularize.regularize_callback.RegularizeCallback._iter_layers': ( 'regularize/regularize_callback.html#regularizecallback._iter_layers',
'fasterai/regularize/regularize_callback.py'),
'fasterai.regularize.regularize_callback.RegularizeCallback.after_epoch': ( 'regularize/regularize_callback.html#regularizecallback.after_epoch',
Expand Down
12 changes: 10 additions & 2 deletions fasterai/regularize/regularize_callback.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,9 +25,12 @@ def __init__(self,
weight: float = 0.01,
layer_types: Type | list[Type] = nn.Conv2d, # Module types to regularize
schedule: Schedule | None = None, # Optional schedule for the weight
verbose: bool = False # Report the weight after each epoch
verbose: bool = False, # Report the weight after each epoch
*,
norm: int = 1, # 1: elementwise L1, 2: group lasso over `granularity`
):
"Callback to apply regularization using criteria during training"
if norm not in (1, 2): raise ValueError(f"norm must be 1 or 2, got {norm!r}")
store_attr()
self.criteria = listify(criteria)
self.granularity = listify(granularity)
Expand All @@ -50,6 +53,11 @@ def _iter_layers(self):
for m in self.learn.model.modules():
if any(isinstance(m, lt) for lt in self.layer_types) and hasattr(m, 'weight'):
yield m

def _group_norms(self, w, dims):
"Per-group L1 (`norm=1`) or L2 (`norm=2`) norms of `w` over `dims`"
if self.norm == 1: return w.abs().sum(dims)
return torch.linalg.vector_norm(w, 2, dim=dims)

def get_norm(self) -> torch.Tensor:
"Compute regularization using the specified criteria and granularities"
Expand All @@ -60,7 +68,7 @@ def get_norm(self) -> torch.Tensor:
for g in self.granularity:
for m in layers:
try:
scores = crit.f(m.weight)[None].abs().sum(Granularities.get_dim(m, g))
scores = self._group_norms(crit.f(m.weight)[None], Granularities.get_dim(m, g))
layer_regs.append(self.current_weight * scores.sum())
except (KeyError, ValueError) as e:
warnings.warn(f"Skipping regularization for {type(m).__name__}: {e}")
Expand Down
68 changes: 32 additions & 36 deletions nbs/regularize/regularize_callback.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@
"source": [
"## Overview\n",
"\n",
"The `RegularizeCallback` applies structured regularization during training to encourage weight sparsity at various granularities. This is useful as a pre-pruning step: by regularizing groups of weights toward zero during training, subsequent pruning can remove more parameters with less accuracy loss.\n",
"The `RegularizeCallback` applies structured regularization during training to encourage weight sparsity at various granularities. This is useful as a pre-pruning step: with `norm=2`, each group of the chosen granularity is penalized as a unit (group lasso), within its layer. In a network where each convolution is followed by BatchNorm, the BatchNorm rescales each filter's output, so penalizing the convolution weights does not select filters; to prepare structured pruning there, penalize the BatchNorm scale or use a group penalty that includes it.\n",
"\n",
"**Key Features:**\n",
"- Supports every granularity defined in `fasterai.core.granularity.Granularities` (for `Conv2d`: `'weight'`, `'shared_weight'`, `'channel'`, `'column'`, `'row'`, `'kernel'`, `'filter'`, and their `shared_`/slice variants)\n",
Expand All @@ -85,9 +85,12 @@
" weight: float = 0.01,\n",
" layer_types: Type | list[Type] = nn.Conv2d, # Module types to regularize\n",
" schedule: Schedule | None = None, # Optional schedule for the weight\n",
" verbose: bool = False # Report the weight after each epoch\n",
" verbose: bool = False, # Report the weight after each epoch\n",
" *,\n",
" norm: int = 1, # 1: elementwise L1, 2: group lasso over `granularity`\n",
" ):\n",
" \"Callback to apply regularization using criteria during training\"\n",
" if norm not in (1, 2): raise ValueError(f\"norm must be 1 or 2, got {norm!r}\")\n",
" store_attr()\n",
" self.criteria = listify(criteria)\n",
" self.granularity = listify(granularity)\n",
Expand All @@ -110,6 +113,11 @@
" for m in self.learn.model.modules():\n",
" if any(isinstance(m, lt) for lt in self.layer_types) and hasattr(m, 'weight'):\n",
" yield m\n",
"\n",
" def _group_norms(self, w, dims):\n",
" \"Per-group L1 (`norm=1`) or L2 (`norm=2`) norms of `w` over `dims`\"\n",
" if self.norm == 1: return w.abs().sum(dims)\n",
" return torch.linalg.vector_norm(w, 2, dim=dims)\n",
" \n",
" def get_norm(self) -> torch.Tensor:\n",
" \"Compute regularization using the specified criteria and granularities\"\n",
Expand All @@ -120,7 +128,7 @@
" for g in self.granularity:\n",
" for m in layers:\n",
" try:\n",
" scores = crit.f(m.weight)[None].abs().sum(Granularities.get_dim(m, g))\n",
" scores = self._group_norms(crit.f(m.weight)[None], Granularities.get_dim(m, g))\n",
" layer_regs.append(self.current_weight * scores.sum())\n",
" except (KeyError, ValueError) as e:\n",
" warnings.warn(f\"Skipping regularization for {type(m).__name__}: {e}\")\n",
Expand All @@ -141,32 +149,24 @@
"id": "6cd300f8",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Found permutation search CUDA kernels\n",
"[ASP][Info] permutation_search_kernels can be imported.\n"
]
},
{
"data": {
"text/markdown": [
"---\n",
"\n",
"[source](https://github.com/FasterAI-Labs/fasterai/tree/master/blob/master/fasterai/regularize/regularize_callback.py#L21){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
"[source](https://github.com/FasterAI-Labs/fasterai/blob/master/fasterai/regularize/regularize_callback.py#L21){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
"\n",
"### RegularizeCallback\n",
"\n",
"```python\n",
"\n",
"def RegularizeCallback(\n",
" criteria:Criteria | list[Criteria], # Criteria(s) to use for regularization\n",
" granularity:str | list[str], # Granularity level(s) for grouping\n",
" weight:float=0.01, # Regularization weight\n",
" layer_types:Type | list[Type]=<class 'torch.nn.modules.conv.Conv2d'>, # Layer types to apply regularization to\n",
" schedule:Schedule | None=None, # Optional schedule for regularization weight\n",
" verbose:bool=False, # Whether to report regularization weight\n",
" criteria:Criteria | list[Criteria], # Importance criteria, e.g. large_final\n",
" granularity:str | list[str], weight:float=0.01,\n",
" layer_types:Type | list[Type]=Conv2d, # Module types to regularize\n",
" schedule:Schedule | None=None, # Optional schedule for the weight\n",
" verbose:bool=False, # Report the weight after each epoch\n",
" norm:int=1, # 1: elementwise L1, 2: group lasso over `granularity`\n",
"):\n",
"\n",
"\n",
Expand All @@ -175,21 +175,15 @@
"*Basic class handling tweaks of the training loop by changing a `Learner` in various events*"
],
"text/plain": [
"```python\n",
"\n",
"def RegularizeCallback(\n",
" criteria:Criteria | list[Criteria], # Criteria(s) to use for regularization\n",
" granularity:str | list[str], # Granularity level(s) for grouping\n",
" weight:float=0.01, # Regularization weight\n",
" layer_types:Type | list[Type]=<class 'torch.nn.modules.conv.Conv2d'>, # Layer types to apply regularization to\n",
" schedule:Schedule | None=None, # Optional schedule for regularization weight\n",
" verbose:bool=False, # Whether to report regularization weight\n",
" criteria:Criteria | list[Criteria], # Importance criteria, e.g. large_final\n",
" granularity:str | list[str], weight:float=0.01,\n",
" layer_types:Type | list[Type]=Conv2d, # Module types to regularize\n",
" schedule:Schedule | None=None, # Optional schedule for the weight\n",
" verbose:bool=False, # Report the weight after each epoch\n",
" norm:int=1, # 1: elementwise L1, 2: group lasso over `granularity`\n",
"):\n",
"\n",
"\n",
"```\n",
"\n",
"*Basic class handling tweaks of the training loop by changing a `Learner` in various events*"
"\"\"\"Basic class handling tweaks of the training loop by changing a `Learner` in various events\"\"\""
]
},
"execution_count": null,
Expand All @@ -208,7 +202,8 @@
"source": [
"**Parameters:**\n",
"- `granularity`: Level at which to group weights, from `Granularities` (e.g. `'weight'`, `'channel'`, `'kernel'`, `'filter'`, `'layer'`)\n",
"- `weight`: Regularization coefficient (higher = stronger regularization)"
"- `weight`: Regularization coefficient (higher = stronger regularization)\n",
"- `norm`: `1` (default) is an elementwise L1, the same value whatever the granularity; `2` sums the L2 norm of each group of the granularity (group lasso), so each group is pushed to zero as a unit. Groups are taken within a single layer: they do not follow cross-layer couplings such as the next layer's input channels or BatchNorm parameters."
]
},
{
Expand All @@ -226,7 +221,7 @@
"source": [
"## Usage Example\n",
"\n",
"Apply filter-level L1 regularization to encourage entire filters to become unimportant (making them easier to prune later):"
"Apply filter-level group lasso (`norm=2`), which penalizes each filter as a unit:"
]
},
{
Expand All @@ -238,11 +233,12 @@
"from fasterai.regularize.regularize_callback import RegularizeCallback\n",
"from fasterai.core.criteria import large_final\n",
"\n",
"# Apply L1 regularization at filter granularity\n",
"# Group lasso at filter granularity: each filter's L2 norm is penalized as a unit\n",
"cb = RegularizeCallback(\n",
" criteria=large_final,\n",
" granularity='filter',\n",
" weight=0.01,\n",
" norm=2,\n",
" verbose=True\n",
")\n",
"\n",
Expand All @@ -256,8 +252,8 @@
"metadata": {},
"source": [
"**Typical Workflow:**\n",
"1. Train with `RegularizeCallback` to push unimportant filter groups toward zero\n",
"2. After training, use `PruneCallback` or `Pruner` to remove the zeroed-out structures\n",
"1. Train with `RegularizeCallback` to penalize groups of weights (in convolution + BatchNorm blocks, see the caveat above)\n",
"2. After training, use `PruneCallback` or `Pruner` to remove the least important structures\n",
"3. Fine-tune the pruned model to recover any lost accuracy"
]
},
Expand Down
97 changes: 92 additions & 5 deletions nbs/tests/test_regularize_callback.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,8 +29,6 @@
"metadata": {},
"outputs": [],
"source": [
"from fastcore.test import *\n",
"\n",
"# Single criteria + granularity\n",
"cb = RegularizeCallback(criteria=large_final, granularity='filter', weight=1e-4)\n",
"test_eq(cb.weight, 1e-4)\n",
Expand All @@ -54,6 +52,93 @@
"test_eq(cb.schedule, None)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "reg-norm1-golden",
"metadata": {},
"outputs": [],
"source": [
"import copy\n",
"from types import SimpleNamespace\n",
"from fastcore.basics import listify\n",
"from fasterai.core.criteria import squared_final\n",
"from fasterai.core.granularity import Granularities\n",
"\n",
"def _reg(model, granularity, norm=1, criteria=large_final, weight=1.):\n",
" cb = RegularizeCallback(criteria, granularity, weight=weight, norm=norm)\n",
" cb.learn = SimpleNamespace(model=model)\n",
" return cb.get_norm()\n",
"\n",
"# norm=1 is byte-identical to the formula before `norm` existed: same loss, same gradients\n",
"def _old_reg(model, crits, grans, weight):\n",
" regs = [weight * crit.f(m.weight)[None].abs().sum(Granularities.get_dim(m, g)).sum()\n",
" for crit in crits for g in grans for m in model.modules() if isinstance(m, nn.Conv2d)]\n",
" return torch.stack(regs).sum()\n",
"\n",
"for g in ['weight', 'filter', 'kernel', ['filter', 'channel']]:\n",
" torch.manual_seed(0); m_new = nn.Sequential(nn.Conv2d(3, 8, 3), nn.ReLU(), nn.Conv2d(8, 4, 3))\n",
" m_old = copy.deepcopy(m_new)\n",
" r_new = _reg(m_new, g, criteria=[large_final, squared_final], weight=0.01); r_new.backward()\n",
" r_old = _old_reg(m_old, [large_final, squared_final], listify(g), 0.01); r_old.backward()\n",
" assert torch.equal(r_new, r_old), g\n",
" for p_new, p_old in zip(m_new.parameters(), m_old.parameters()):\n",
" assert (p_new.grad is None and p_old.grad is None) or torch.equal(p_new.grad, p_old.grad), g"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "reg-granularity",
"metadata": {},
"outputs": [],
"source": [
"# norm=1 is the same value whatever the granularity; norm=2 is not\n",
"torch.manual_seed(0); _m = nn.Conv2d(3, 8, 3)\n",
"test_close(_reg(_m, 'filter'), _reg(_m, 'weight'), eps=1e-4)\n",
"test_close(_reg(_m, 'kernel'), _reg(_m, 'weight'), eps=1e-4)\n",
"test_ne(_reg(_m, 'filter', norm=2).item(), _reg(_m, 'kernel', norm=2).item())"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "reg-norm2-values",
"metadata": {},
"outputs": [],
"source": [
"def _conv(*filters):\n",
" \"Conv2d(1, len(filters), 2) whose filters are the given 4-element lists\"\n",
" m = nn.Conv2d(1, len(filters), 2, bias=False)\n",
" with torch.no_grad(): m.weight.copy_(torch.tensor(filters, dtype=torch.float).view(-1, 1, 2, 2))\n",
" return m\n",
"\n",
"# by hand: filters [3,4,0,0] and [1,1,1,1] have L2 norms 5 and 2, L1 norms 7 and 4\n",
"_c = _conv([3, 4, 0, 0], [1, 1, 1, 1])\n",
"test_close(_reg(_c, 'filter', norm=2), 7., eps=1e-5)\n",
"test_close(_reg(_c, 'filter', norm=2, weight=0.1), 0.7, eps=1e-6)\n",
"test_close(_reg(_c, 'weight', norm=2), 11., eps=1e-5) # singleton groups: L1\n",
"\n",
"# gradient is finite (exactly zero) on a group that is exactly zero\n",
"_z = _conv([0, 0, 0, 0], [1, 1, 1, 1])\n",
"_reg(_z, 'filter', norm=2).backward()\n",
"assert torch.isfinite(_z.weight.grad).all()\n",
"test_eq(_z.weight.grad[0].abs().sum().item(), 0.)\n",
"test_close(_z.weight.grad[1].flatten(), torch.full((4,), 0.5), eps=1e-5)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "reg-bad-norm",
"metadata": {},
"outputs": [],
"source": [
"# any other norm is refused\n",
"for bad in (0, 3, 'l2', 2.5):\n",
" with ExceptionExpected(ValueError, 'norm must be 1 or 2'): RegularizeCallback(large_final, 'filter', norm=bad)"
]
},
{
"cell_type": "code",
"execution_count": null,
Expand All @@ -80,9 +165,11 @@
" bs=16, device='cpu'\n",
")\n",
"\n",
"_cb = RegularizeCallback(criteria=large_final, granularity='filter', weight=1e-4)\n",
"_learn = Learner(_dls, _model, loss_func=nn.CrossEntropyLoss(), cbs=[_cb])\n",
"_learn.fit(2) # verify it runs end-to-end without error"
"for _norm in (1, 2):\n",
" _cb = RegularizeCallback(criteria=large_final, granularity='filter', weight=1e-4, norm=_norm)\n",
" _learn = Learner(_dls, _model, loss_func=nn.CrossEntropyLoss(), cbs=[_cb])\n",
" _learn.fit(2) # verify it runs end-to-end without error\n",
" assert all(torch.isfinite(p).all() for p in _model.parameters())"
]
}
],
Expand Down
Loading
Loading