From bab7b3e81a7fbe20232bd972c81c238325807a7d Mon Sep 17 00:00:00 2001 From: nathanhubens Date: Thu, 24 Sep 2026 12:41:51 +0200 Subject: [PATCH] GroupRegularizeCallback: apply a Pruner's group penalty during a fastai fit before_step calls pruner.regularize(), ordered between GradientAccumulation and MixedPrecision, so the penalty is added once per optimizer step and multiplied by the loss scale of an active torch.amp.GradScaler. An optional Schedule ramps it and is copied, leaving the shared module-level schedules untouched. Pruning stays a separate call after the fit, followed by learn.fit(n, reset_opt=True). A Pruner with reg=0, a Pruner built on another model, NonNativeMixedPrecision and a PruneCallback in the same fit are refused; PruneCallback refuses reg and points to the new callback. --- CHANGELOG.md | 3 +- fasterai/_modidx.py | 10 + fasterai/prune/prune_callback.py | 1 + fasterai/regularize/all.py | 3 +- .../regularize/group_regularize_callback.py | 42 +++ nbs/_quarto.yml | 1 + nbs/prune/prune_callback.ipynb | 2 + nbs/prune/pruner.ipynb | 3 +- .../group_regularize_callback.ipynb | 197 +++++++++++++ nbs/regularize/regularize_callback.ipynb | 3 +- .../test_group_regularize_callback.ipynb | 261 ++++++++++++++++++ nbs/tests/test_prune_callback.ipynb | 12 + 12 files changed, 534 insertions(+), 4 deletions(-) create mode 100644 fasterai/regularize/group_regularize_callback.py create mode 100644 nbs/regularize/group_regularize_callback.ipynb create mode 100644 nbs/tests/test_group_regularize_callback.ipynb diff --git a/CHANGELOG.md b/CHANGELOG.md index 655906ea..81124381 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,8 @@ ## Unreleased ### 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 +- `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` refuses `reg`: use `GroupRegularizeCallback` +- `GroupRegularizeCallback(pruner, schedule=None)` applies a `Pruner`'s group penalty during a fastai fit: `before_step` calls `pruner.regularize()`, after `GradientAccumulation` has skipped the batches that do not step and before `MixedPrecision` unscales the gradients, so the penalty is added once per optimizer step and multiplied by the loss scale of a `torch.amp.GradScaler` when one is active. A `Schedule` multiplies the penalty by its progress and is copied, so the module-level schedules other callbacks share are not advanced. Pruning stays a separate call after the fit: `pruner.prune_model()`, then `learn.fit(n, reset_opt=True)`, because the prune replaces the parameters the optimizer holds. A `Pruner` with `reg=0` is refused, and so are, at `before_fit`, a `Pruner` built on another model than `learn.model`, `NonNativeMixedPrecision` and a `PruneCallback` in the same fit - `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 ### Breaking Changes diff --git a/fasterai/_modidx.py b/fasterai/_modidx.py index 6e8d6471..89212ede 100644 --- a/fasterai/_modidx.py +++ b/fasterai/_modidx.py @@ -636,6 +636,16 @@ 'fasterai/quantize/quantizer.py'), 'fasterai.quantize.quantizer._torchao_recipe': ( 'quantize/quantizer.html#_torchao_recipe', 'fasterai/quantize/quantizer.py')}, + 'fasterai.regularize.group_regularize_callback': { 'fasterai.regularize.group_regularize_callback.GroupRegularizeCallback': ( 'regularize/group_regularize_callback.html#groupregularizecallback', + 'fasterai/regularize/group_regularize_callback.py'), + 'fasterai.regularize.group_regularize_callback.GroupRegularizeCallback.__init__': ( 'regularize/group_regularize_callback.html#groupregularizecallback.__init__', + 'fasterai/regularize/group_regularize_callback.py'), + 'fasterai.regularize.group_regularize_callback.GroupRegularizeCallback.after_epoch': ( 'regularize/group_regularize_callback.html#groupregularizecallback.after_epoch', + 'fasterai/regularize/group_regularize_callback.py'), + 'fasterai.regularize.group_regularize_callback.GroupRegularizeCallback.before_fit': ( 'regularize/group_regularize_callback.html#groupregularizecallback.before_fit', + 'fasterai/regularize/group_regularize_callback.py'), + 'fasterai.regularize.group_regularize_callback.GroupRegularizeCallback.before_step': ( 'regularize/group_regularize_callback.html#groupregularizecallback.before_step', + 'fasterai/regularize/group_regularize_callback.py')}, 'fasterai.regularize.regularize_callback': { 'fasterai.regularize.regularize_callback.RegularizeCallback': ( 'regularize/regularize_callback.html#regularizecallback', 'fasterai/regularize/regularize_callback.py'), 'fasterai.regularize.regularize_callback.RegularizeCallback.__init__': ( 'regularize/regularize_callback.html#regularizecallback.__init__', diff --git a/fasterai/prune/prune_callback.py b/fasterai/prune/prune_callback.py index cd64bd82..e7157017 100644 --- a/fasterai/prune/prune_callback.py +++ b/fasterai/prune/prune_callback.py @@ -34,6 +34,7 @@ def __init__(self, *args, **kwargs ): + if 'reg' in kwargs: raise ValueError("PruneCallback does not regularize: give reg to a Pruner and train with GroupRegularizeCallback.") store_attr() self.sparsity_levels = [] self._is_per_layer = False diff --git a/fasterai/regularize/all.py b/fasterai/regularize/all.py index 1b286145..540955ef 100644 --- a/fasterai/regularize/all.py +++ b/fasterai/regularize/all.py @@ -1 +1,2 @@ -from .regularize_callback import * \ No newline at end of file +from .regularize_callback import * +from .group_regularize_callback import * diff --git a/fasterai/regularize/group_regularize_callback.py b/fasterai/regularize/group_regularize_callback.py new file mode 100644 index 00000000..5b61ba22 --- /dev/null +++ b/fasterai/regularize/group_regularize_callback.py @@ -0,0 +1,42 @@ +# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/regularize/group_regularize_callback.ipynb. + +# %% ../../nbs/regularize/group_regularize_callback.ipynb #imports +from __future__ import annotations +from copy import copy + +from fastai.callback.core import Callback +from fastai.callback.fp16 import MixedPrecision, NonNativeMixedPrecision +from torch.amp import GradScaler +from ..core.schedule import Schedule +from ..prune.pruner import Pruner +from ..prune.prune_callback import PruneCallback + +# %% auto #0 +__all__ = ['GroupRegularizeCallback'] + +# %% ../../nbs/regularize/group_regularize_callback.ipynb #group-reg-cb +class GroupRegularizeCallback(Callback): + "Add `pruner`'s group penalty to the gradients at every optimizer step, so training pushes coupled channels towards zero before `pruner.prune_model()` removes them" + order = MixedPrecision.order - 2 # after GradientAccumulation skips a batch, before MixedPrecision unscales + def __init__(self, + pruner: Pruner, # Built with `reg > 0` on the model the Learner trains + schedule: Schedule | None = None, # Ramps the penalty: `pruner.reg * progress` + verbose: bool = False, # Print the penalty after each epoch + ): + if not pruner.reg: raise ValueError("GroupRegularizeCallback needs a penalty: build the Pruner with reg > 0.") + self.pruner, self.schedule, self.verbose, self.scale = pruner, copy(schedule), verbose, 1. + + def before_fit(self): + if self.pruner.model is not self.learn.model: raise ValueError("The Pruner was built on another model than learn.model: build it on learn.model.") + if any(isinstance(cb, NonNativeMixedPrecision) for cb in self.learn.cbs): raise ValueError("NonNativeMixedPrecision is not supported: use MixedPrecision (learn.to_fp16()).") + if any(isinstance(cb, PruneCallback) for cb in self.learn.cbs): raise ValueError("PruneCallback cannot run in the same fit: fit with GroupRegularizeCallback, then call pruner.prune_model().") + + def before_step(self): + self.scale = self.schedule.progress(self.pct_train) if self.schedule is not None else 1. + if not self.scale: return + scaler = getattr(self.learn, 'scaler', None) + loss_scale = scaler.get_scale() if isinstance(scaler, GradScaler) else 1. + self.pruner.regularize(scale=self.scale * loss_scale) + + def after_epoch(self): + if self.verbose: print(f"Group penalty: {self.pruner.reg * self.scale:.2e}") diff --git a/nbs/_quarto.yml b/nbs/_quarto.yml index da9fedf9..07584a4b 100644 --- a/nbs/_quarto.yml +++ b/nbs/_quarto.yml @@ -117,6 +117,7 @@ website: - section: Regularize contents: - regularize/regularize_callback.ipynb + - regularize/group_regularize_callback.ipynb - section: Misc contents: - misc/bn_folding.ipynb diff --git a/nbs/prune/prune_callback.ipynb b/nbs/prune/prune_callback.ipynb index 265b47aa..38e6e8e0 100644 --- a/nbs/prune/prune_callback.ipynb +++ b/nbs/prune/prune_callback.ipynb @@ -94,6 +94,7 @@ " *args,\n", " **kwargs\n", " ):\n", + " if 'reg' in kwargs: raise ValueError(\"PruneCallback does not regularize: give reg to a Pruner and train with GroupRegularizeCallback.\")\n", " store_attr()\n", " self.sparsity_levels = []\n", " self._is_per_layer = False\n", @@ -240,6 +241,7 @@ "## See Also\n", "\n", "- [Pruner](pruner.html) - Core structured pruning class used by this callback\n", + "- [GroupRegularizeCallback](../regularize/group_regularize_callback.html) - Train with a `Pruner`'s group penalty, then prune after the fit (`PruneCallback` does not take `reg`)\n", "- [Schedules](../core/schedules.html) - Control pruning progression during training\n", "- [Criteria](../core/criteria.html) - Importance measures for selecting filters to prune\n", "- [SparsifyCallback](../sparse/sparsify_callback.html) - Unstructured pruning alternative\n", diff --git a/nbs/prune/pruner.ipynb b/nbs/prune/pruner.ipynb index 8c922ddc..3ed1d33d 100644 --- a/nbs/prune/pruner.ipynb +++ b/nbs/prune/pruner.ipynb @@ -573,7 +573,7 @@ "pruner.prune_model()\n", "```\n", "\n", - "Every channel of a 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` goes from `1` for the most important channel of the group to `2**alpha` for the least important, on the square root of the `criteria` score. Only the groups the pruner will prune are penalized (a per-layer dict leaves the others alone), biases and frozen parameters are not, a depthwise convolution's own kernels are neither scored nor penalized, and a group whose importance is constant is skipped. The groups are rebuilt after every `prune_model()`. With `reg=0` (the default) nothing changes. If the gradients are scaled (a mixed-precision loss scale, gradient accumulation), pass that factor as `scale`." + "Every channel of a 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` goes from `1` for the most important channel of the group to `2**alpha` for the least important, on the square root of the `criteria` score. Only the groups the pruner will prune are penalized (a per-layer dict leaves the others alone), biases and frozen parameters are not, a depthwise convolution's own kernels are neither scored nor penalized, and a group whose importance is constant is skipped. The groups are rebuilt after every `prune_model()`. With `reg=0` (the default) nothing changes. If the gradients are scaled (a mixed-precision loss scale, gradient accumulation), pass that factor as `scale`. In a fastai `Learner`, [GroupRegularizeCallback](../regularize/group_regularize_callback.html) makes this call at every optimizer step." ] }, { @@ -586,6 +586,7 @@ "## See Also\n", "\n", "- [PruneCallback](prune_callback.html) - Apply structured pruning during fastai training\n", + "- [GroupRegularizeCallback](../regularize/group_regularize_callback.html) - Apply the group penalty during a fastai fit\n", "- [Criteria](../core/criteria.html) - Different importance measures for selecting what to prune\n", "- [Schedules](../core/schedules.html) - Control pruning progression during training\n", "- [Sparsifier](../sparse/sparsifier.html) - Unstructured pruning (zeroing weights without removing them)\n", diff --git a/nbs/regularize/group_regularize_callback.ipynb b/nbs/regularize/group_regularize_callback.ipynb new file mode 100644 index 00000000..44883b77 --- /dev/null +++ b/nbs/regularize/group_regularize_callback.ipynb @@ -0,0 +1,197 @@ +{ + "cells": [ + { + "cell_type": "raw", + "id": "frontmatter", + "metadata": {}, + "source": [ + "---\n", + "description: Train with a Pruner's group penalty before pruning\n", + "output-file: group_regularize_callback.html\n", + "title: Group Regularize Callback\n", + "---\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "default-exp", + "metadata": {}, + "outputs": [], + "source": [ + "#| default_exp regularize.group_regularize_callback" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "showdoc-import", + "metadata": {}, + "outputs": [], + "source": [ + "#| include: false\n", + "from nbdev.showdoc import *" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "imports", + "metadata": {}, + "outputs": [], + "source": [ + "#| export\n", + "from __future__ import annotations\n", + "from copy import copy\n", + "\n", + "from fastai.callback.core import Callback\n", + "from fastai.callback.fp16 import MixedPrecision, NonNativeMixedPrecision\n", + "from torch.amp import GradScaler\n", + "from fasterai.core.schedule import Schedule\n", + "from fasterai.prune.pruner import Pruner\n", + "from fasterai.prune.prune_callback import PruneCallback" + ] + }, + { + "cell_type": "markdown", + "id": "overview", + "metadata": {}, + "source": [ + "## Overview\n", + "\n", + "`GroupRegularizeCallback` calls `Pruner.regularize()` at every optimizer step of a fastai fit: the gradients receive the [group penalty](../prune/pruner.html#group-sparse-learning) of a `Pruner` built with `reg > 0`, so training pushes the coupled channels the pruner will remove (a convolution's filters, its BatchNorm scale, the input slices of the layers that read it) towards zero. The same `Pruner` then removes them with `prune_model()`, once the fit is over." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "group-reg-cb", + "metadata": {}, + "outputs": [], + "source": [ + "#| export\n", + "class GroupRegularizeCallback(Callback):\n", + " \"Add `pruner`'s group penalty to the gradients at every optimizer step, so training pushes coupled channels towards zero before `pruner.prune_model()` removes them\"\n", + " order = MixedPrecision.order - 2 # after GradientAccumulation skips a batch, before MixedPrecision unscales\n", + " def __init__(self,\n", + " pruner: Pruner, # Built with `reg > 0` on the model the Learner trains\n", + " schedule: Schedule | None = None, # Ramps the penalty: `pruner.reg * progress`\n", + " verbose: bool = False, # Print the penalty after each epoch\n", + " ):\n", + " if not pruner.reg: raise ValueError(\"GroupRegularizeCallback needs a penalty: build the Pruner with reg > 0.\")\n", + " self.pruner, self.schedule, self.verbose, self.scale = pruner, copy(schedule), verbose, 1.\n", + "\n", + " def before_fit(self):\n", + " if self.pruner.model is not self.learn.model: raise ValueError(\"The Pruner was built on another model than learn.model: build it on learn.model.\")\n", + " if any(isinstance(cb, NonNativeMixedPrecision) for cb in self.learn.cbs): raise ValueError(\"NonNativeMixedPrecision is not supported: use MixedPrecision (learn.to_fp16()).\")\n", + " if any(isinstance(cb, PruneCallback) for cb in self.learn.cbs): raise ValueError(\"PruneCallback cannot run in the same fit: fit with GroupRegularizeCallback, then call pruner.prune_model().\")\n", + "\n", + " def before_step(self):\n", + " self.scale = self.schedule.progress(self.pct_train) if self.schedule is not None else 1.\n", + " if not self.scale: return\n", + " scaler = getattr(self.learn, 'scaler', None)\n", + " loss_scale = scaler.get_scale() if isinstance(scaler, GradScaler) else 1.\n", + " self.pruner.regularize(scale=self.scale * loss_scale)\n", + "\n", + " def after_epoch(self):\n", + " if self.verbose: print(f\"Group penalty: {self.pruner.reg * self.scale:.2e}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "show-doc", + "metadata": {}, + "outputs": [ + { + "data": { + "text/markdown": [ + "---\n", + "\n", + "[source](https://github.com/FasterAI-Labs/fasterai/blob/master/fasterai/regularize/group_regularize_callback.py#L18){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n", + "\n", + "### GroupRegularizeCallback\n", + "\n", + "```python\n", + "def GroupRegularizeCallback(\n", + " pruner:Pruner, # Built with `reg > 0` on the model the Learner trains\n", + " schedule:Schedule | None=None, # Ramps the penalty: `pruner.reg * progress`\n", + " verbose:bool=False, # Print the penalty after each epoch\n", + "):\n", + "```\n", + "\n", + "*Add `pruner`'s group penalty to the gradients at every optimizer step, so training pushes coupled channels towards zero before `pruner.prune_model()` removes them*" + ], + "text/plain": [ + "def GroupRegularizeCallback(\n", + " pruner:Pruner, # Built with `reg > 0` on the model the Learner trains\n", + " schedule:Schedule | None=None, # Ramps the penalty: `pruner.reg * progress`\n", + " verbose:bool=False, # Print the penalty after each epoch\n", + "):\"\"\"Add `pruner`'s group penalty to the gradients at every optimizer step, so training pushes coupled channels towards zero before `pruner.prune_model()` removes them\"\"\"" + ] + }, + "execution_count": null, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "show_doc(GroupRegularizeCallback)" + ] + }, + { + "cell_type": "markdown", + "id": "usage", + "metadata": {}, + "source": [ + "## Usage Example\n", + "\n", + "Build the `Pruner` on `learn.model` with `reg > 0`, train with the callback, prune, then fine-tune:\n", + "\n", + "```python\n", + "from fasterai.core.criteria import large_final\n", + "from fasterai.core.schedule import lin\n", + "from fasterai.prune.pruner import Pruner\n", + "from fasterai.regularize.group_regularize_callback import GroupRegularizeCallback\n", + "\n", + "xb, _ = learn.dls.one_batch()\n", + "pruner = Pruner(learn.model, 0.3, 'local', large_final, reg=1e-4, example_inputs=xb)\n", + "learn.fit(10, cbs=GroupRegularizeCallback(pruner, schedule=lin))\n", + "pruner.prune_model()\n", + "learn.fit(2, reset_opt=True)\n", + "```\n", + "\n", + "- The prune replaces the parameters of every layer it shrinks, so the next fit needs `reset_opt=True`: without it, fastai keeps the optimizer it built on the old parameters.\n", + "- It works under `learn.to_fp16()` and `GradientAccumulation`: the penalty is multiplied by the loss scale before `MixedPrecision` unscales the gradients, and added once per optimizer step, not once per batch. `NonNativeMixedPrecision` is refused.\n", + "- Frozen layers get no penalty.\n", + "- Prune after the fit, not during: `PruneCallback`, which prunes during training, is refused in the same fit." + ] + }, + { + "cell_type": "markdown", + "id": "see-also", + "metadata": {}, + "source": [ + "---\n", + "\n", + "## See Also\n", + "\n", + "- [Pruner](../prune/pruner.html) - `reg`, `alpha` and `regularize()`, the penalty this callback applies\n", + "- [PruneCallback](../prune/prune_callback.html) - Prune during training instead of after it\n", + "- [RegularizeCallback](regularize_callback.html) - Penalize each layer's own weights, one layer at a time\n", + "- [Schedules](../core/schedules.html) - Ramp the penalty over training\n", + "\n", + "Tests live in `nbs/tests/test_group_regularize_callback.ipynb`." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "python3", + "language": "python", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/nbs/regularize/regularize_callback.ipynb b/nbs/regularize/regularize_callback.ipynb index 30f79355..2d1c6538 100644 --- a/nbs/regularize/regularize_callback.ipynb +++ b/nbs/regularize/regularize_callback.ipynb @@ -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: 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", + "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 [GroupRegularizeCallback](group_regularize_callback.html), whose group penalty 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", @@ -266,6 +266,7 @@ "\n", "## See Also\n", "\n", + "- [GroupRegularizeCallback](group_regularize_callback.html) - Penalize coupled channels across layers, BatchNorm included, before a `Pruner` removes them\n", "- [Sparsifier](../sparse/sparsifier.html) - Apply sparsification after regularization pushes weights to zero\n", "- [Criteria](../core/criteria.html) - Importance measures that can leverage regularized weights\n", "- [SparsifyCallback](../sparse/sparsify_callback.html) - Combine with sparsification for gradual pruning\n", diff --git a/nbs/tests/test_group_regularize_callback.ipynb b/nbs/tests/test_group_regularize_callback.ipynb new file mode 100644 index 00000000..621d8859 --- /dev/null +++ b/nbs/tests/test_group_regularize_callback.ipynb @@ -0,0 +1,261 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "title", + "metadata": {}, + "source": [ + "# Tests: `fasterai.regularize.group_regularize_callback` (source `nbs/regularize/group_regularize_callback.ipynb`)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "imports", + "metadata": {}, + "outputs": [], + "source": [ + "from fastcore.test import *\n", + "from copy import deepcopy\n", + "import tempfile\n", + "import torch\n", + "import torch.nn as nn\n", + "from torch.amp import GradScaler\n", + "from torch.utils.data import TensorDataset\n", + "from fastai.callback.core import Callback\n", + "from fastai.callback.fp16 import MixedPrecision, NonNativeMixedPrecision\n", + "from fastai.callback.training import GradientAccumulation\n", + "from fastai.data.core import DataLoaders\n", + "from fastai.learner import Learner\n", + "from fastai.optimizer import SGD\n", + "from fasterai.core.criteria import large_final\n", + "from fasterai.core.schedule import lin, one_shot\n", + "from fasterai.prune.pruner import Pruner\n", + "from fasterai.prune.prune_callback import PruneCallback\n", + "from fasterai.regularize.group_regularize_callback import *" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "penalty-by-hand", + "metadata": {}, + "outputs": [], + "source": [ + "# --- the penalty reaches the gradients through a real fit step ---\n", + "# values from test_pruner's `_Tiny`, reg 0.01; a zero data loss and lr 0 leave only the penalty\n", + "class _Tiny(nn.Module):\n", + " def __init__(self):\n", + " super().__init__()\n", + " self.conv, self.bn, self.fc = nn.Conv2d(1, 2, 1, bias=False), nn.BatchNorm2d(2), nn.Linear(2, 1)\n", + " with torch.no_grad(): self.conv.weight.copy_(torch.tensor([1., 4.]).view(2, 1, 1, 1)); self.fc.weight.fill_(1.)\n", + " def forward(self, x): return self.fc(self.bn(self.conv(x)).mean((2, 3)))\n", + "\n", + "_EXPECTED = {'conv.weight': torch.tensor([0.08, 0.04]), 'bn.weight': torch.tensor([0.08, 0.01]), 'fc.weight': torch.tensor([0.08, 0.01])}\n", + "\n", + "class _Probe(Callback):\n", + " \"Record the gradients the optimizer is about to see, at `order`\"\n", + " def __init__(self, order=60): self.order, self.grads = order, []\n", + " def before_step(self): self.grads.append({n: p.grad.detach().flatten().clone() for n, p in self.learn.model.named_parameters()})\n", + "\n", + "def _fit(learn, *cbs, n=1, **kw):\n", + " with learn.no_bar(), learn.no_logging(): learn.fit(n, cbs=list(cbs), **kw)\n", + "\n", + "def _tiny_learner(n=8, bs=2, reg=0.01):\n", + " m = _Tiny()\n", + " X, y = torch.randn(n, 1, 2, 2), torch.zeros(n, 1)\n", + " dls = DataLoaders.from_dsets(TensorDataset(X, y), TensorDataset(X[:bs], y[:bs]), bs=bs, device='cpu')\n", + " learn = Learner(dls, m, loss_func=lambda p, y: p.sum() * 0, opt_func=SGD, lr=0.)\n", + " return learn, Pruner(m, 0.5, 'local', large_final, example_inputs=X[:bs], reg=reg, alpha=3)\n", + "\n", + "def _run(*cbs, schedule=None, probe=None, **learner_kw):\n", + " \"Fit a fresh `_Tiny` learner with the callback, `cbs` and a probe; return the probe\"\n", + " learn, pr = _tiny_learner(**learner_kw)\n", + " probe = probe or _Probe()\n", + " _fit(learn, GroupRegularizeCallback(pr, schedule=schedule), *cbs, probe)\n", + " return probe\n", + "\n", + "def _check(g, factor=1.):\n", + " for n, e in _EXPECTED.items(): test_close(g[n], factor * e, eps=1e-6)\n", + " test_eq(g['bn.bias'].abs().max().item(), 0.); test_eq(g['fc.bias'].abs().max().item(), 0.)\n", + "\n", + "probe = _run()\n", + "test_eq(len(probe.grads), 4)\n", + "for g in probe.grads: _check(g)\n", + "\n", + "# without the callback, the same fit sees zero gradients\n", + "learn, pr = _tiny_learner(); probe = _Probe()\n", + "_fit(learn, probe)\n", + "assert all(t.abs().max() == 0 for t in probe.grads[0].values())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "order", + "metadata": {}, + "outputs": [], + "source": [ + "# --- order: strictly after GradientAccumulation, strictly before MixedPrecision (which unscales) ---\n", + "assert GradientAccumulation.order < GroupRegularizeCallback.order < MixedPrecision.order" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "amp", + "metadata": {}, + "outputs": [], + "source": [ + "# --- mixed precision: a scaled gradient gets a scaled penalty, the optimizer sees the true one ---\n", + "# CPU stub of MixedPrecision (order 10): a real GradScaler scales the loss, the gradients are divided back at the step\n", + "S = 1024.\n", + "class _Scaled(Callback):\n", + " order = MixedPrecision.order\n", + " def before_fit(self): self.learn.scaler = GradScaler('cpu', init_scale=S)\n", + " def before_backward(self): self.learn.loss_grad = self.learn.scaler.scale(self.learn.loss_grad)\n", + " def before_step(self):\n", + " for p in self.learn.model.parameters(): p.grad /= S\n", + " def after_fit(self): self.learn.scaler = None\n", + "\n", + "before = _Probe(order=MixedPrecision.order - 1)\n", + "after = _run(_Scaled(), before)\n", + "for gb, ga in zip(before.grads, after.grads): _check(gb, S); _check(ga)\n", + "\n", + "# only a real GradScaler is read: `learn.scaler` falls through to the model, whose attribute is ignored\n", + "learn, pr = _tiny_learner()\n", + "learn.model.scaler = 7.\n", + "probe = _Probe()\n", + "_fit(learn, GroupRegularizeCallback(pr), probe)\n", + "_check(probe.grads[0])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "grad-accum", + "metadata": {}, + "outputs": [], + "source": [ + "# --- gradient accumulation: one penalty per real optimizer step, not per batch ---\n", + "learn, pr = _tiny_learner(n=8, bs=2)\n", + "calls, _reg = [], pr.regularize\n", + "pr.regularize = lambda scale=1.: (calls.append(scale), _reg(scale))\n", + "probe = _Probe()\n", + "_fit(learn, GroupRegularizeCallback(pr), GradientAccumulation(n_acc=4), probe)\n", + "test_eq(len(calls), 2); test_eq(len(probe.grads), 2)\n", + "for g in probe.grads: _check(g)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "schedule", + "metadata": {}, + "outputs": [], + "source": [ + "# --- a schedule ramps the penalty, and the shared module-level schedule is left alone ---\n", + "for sched, progress in ((one_shot, [0, 0, 1, 1]), (lin, [0, .25, .5, .75])):\n", + " state = deepcopy(sched.__dict__)\n", + " probe = _run(schedule=sched)\n", + " for g, f in zip(probe.grads, progress): _check(g, f)\n", + " test_eq(sched.__dict__, state)\n", + "assert GroupRegularizeCallback(_tiny_learner()[1], schedule=lin).schedule is not lin\n", + "\n", + "# verbose prints the penalty reached at the end of each epoch\n", + "learn, pr = _tiny_learner()\n", + "test_stdout(lambda: _fit(learn, GroupRegularizeCallback(pr, schedule=lin, verbose=True)), 'Group penalty: 7.50e-03')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "refusals", + "metadata": {}, + "outputs": [], + "source": [ + "# --- refusals, each with one line naming the fix ---\n", + "learn, pr = _tiny_learner()\n", + "with ExceptionExpected(ValueError, regex='reg > 0'): GroupRegularizeCallback(_tiny_learner(reg=0.)[1])\n", + "\n", + "_other = _tiny_learner()[1] # a Pruner built on another model\n", + "with ExceptionExpected(ValueError, regex='another model than learn.model'): _fit(learn, GroupRegularizeCallback(_other))\n", + "with ExceptionExpected(ValueError, regex='use MixedPrecision'): _fit(learn, GroupRegularizeCallback(pr), NonNativeMixedPrecision())\n", + "with ExceptionExpected(ValueError, regex='PruneCallback cannot run in the same fit'):\n", + " _fit(learn, GroupRegularizeCallback(pr), PruneCallback(0.5, one_shot, 'local', large_final, example_inputs=torch.randn(2, 1, 2, 2)))\n", + "\n", + "# a refused fit leaves no callback behind: the next fit runs\n", + "_fit(learn, GroupRegularizeCallback(pr))\n", + "assert not any(isinstance(c, GroupRegularizeCallback) for c in learn.cbs)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "protocol", + "metadata": {}, + "outputs": [], + "source": [ + "# --- the documented protocol: fit with the callback, prune after it, fine-tune with reset_opt=True ---\n", + "_live = lambda model: {id(p) for p in model.parameters()}\n", + "_held = lambda learn: {id(p) for g in learn.opt.param_lists for p in g}\n", + "\n", + "torch.manual_seed(0)\n", + "_model = nn.Sequential(nn.Conv2d(3, 8, 3, padding=1), nn.BatchNorm2d(8), nn.ReLU(),\n", + " nn.Conv2d(8, 16, 3, padding=1), nn.BatchNorm2d(16), nn.ReLU(),\n", + " nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(16, 4))\n", + "_X, _y = torch.randn(32, 3, 8, 8), torch.randint(0, 4, (32,))\n", + "_dls = DataLoaders.from_dsets(TensorDataset(_X[:24], _y[:24]), TensorDataset(_X[24:], _y[24:]), bs=8, device='cpu')\n", + "_learn = Learner(_dls, _model, loss_func=nn.CrossEntropyLoss(), lr=1e-2)\n", + "_pr = Pruner(_model, 0.5, 'local', large_final, example_inputs=_X[:2], reg=1e-3)\n", + "_fit(_learn, GroupRegularizeCallback(_pr))\n", + "_pr.prune_model()\n", + "test_eq(_model[0].out_channels, 4)\n", + "\n", + "assert _held(_learn) != _live(_model) # why reset_opt=True: the prune replaced the parameters fastai's optimizer holds\n", + "w0 = _model[0].weight.detach().clone()\n", + "_fit(_learn, reset_opt=True)\n", + "test_eq(_held(_learn), _live(_model))\n", + "assert not torch.equal(_model[0].weight, w0)\n", + "test_eq(_model(_X[:2]).shape, (2, 4))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "slow-resnet18", + "metadata": {}, + "outputs": [], + "source": [ + "#| slow\n", + "# --- resnet18 on the CPU: fit with the callback, prune, fine-tune, save ---\n", + "from torchvision.models import resnet18\n", + "torch.manual_seed(0)\n", + "_rm = resnet18(num_classes=4)\n", + "_X, _y = torch.randn(32, 3, 32, 32), torch.randint(0, 4, (32,))\n", + "_dls = DataLoaders.from_dsets(TensorDataset(_X[:24], _y[:24]), TensorDataset(_X[24:], _y[24:]), bs=8, device='cpu')\n", + "_learn = Learner(_dls, _rm, loss_func=nn.CrossEntropyLoss(), lr=1e-3, path=tempfile.mkdtemp())\n", + "_pr = Pruner(_rm, 0.3, 'local', large_final, example_inputs=_X[:2], reg=1e-4)\n", + "_probe = _Probe()\n", + "_fit(_learn, GroupRegularizeCallback(_pr, schedule=lin), _probe)\n", + "assert all(t.isfinite().all() for g in _probe.grads for t in g.values())\n", + "n0 = sum(p.numel() for p in _rm.parameters())\n", + "_pr.prune_model()\n", + "assert sum(p.numel() for p in _rm.parameters()) < n0\n", + "_fit(_learn, reset_opt=True)\n", + "test_eq(_held(_learn), _live(_rm))\n", + "test_eq(_rm(_X[:2]).shape, (2, 4))\n", + "assert _learn.save('pruned').exists()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "python3", + "language": "python", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/nbs/tests/test_prune_callback.ipynb b/nbs/tests/test_prune_callback.ipynb index 0f32d4e7..2d79c8cd 100644 --- a/nbs/tests/test_prune_callback.ipynb +++ b/nbs/tests/test_prune_callback.ipynb @@ -94,6 +94,18 @@ " PruneCallback(pruning_ratio='0.3', schedule=one_shot, context='local', criteria=large_final)" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "reg-refused", + "metadata": {}, + "outputs": [], + "source": [ + "# `reg` is refused by name: the group penalty is GroupRegularizeCallback's job\n", + "with ExceptionExpected(ValueError, regex='GroupRegularizeCallback'):\n", + " PruneCallback(0.3, agp, 'local', large_final, reg=1e-4)" + ] + }, { "cell_type": "code", "execution_count": null,