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,