diff --git a/CHANGELOG.md b/CHANGELOG.md index ab99f7be..655906ea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,9 @@ - `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 +### Breaking Changes +- `PruneCallback` and `Pruner(schedule=Schedule)` now honour the schedule's `start_pct` and `end_pct`. The levels handed to torch-pruning used to map step `i` of `n` straight into the schedule's shape function, so `start_pct` and `end_pct` were never read: `one_shot` (`start_pct=0.5`) pruned at the first step, and `agp` and `iterative` (`start_pct=0.2`) began ramping from the first step. The levels are now the schedule's own `progress(i / n)`, built in one place (`Pruner._to_tp_scheduler`) for both entry points: `one_shot` prunes once, at the first step at or past `start_pct`; `agp` and `iterative` begin at their `start_pct`; the full ratio is reached at `end_pct` and held after it. A schedule with `start_pct=0` and `end_pct=1` (`one_cycle`, `cos`, `lin`, `Schedule(f)`) produces the same levels as before, level for level, and a plain torch-pruning scheduler callable is passed through unchanged + ### 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 diff --git a/fasterai/_modidx.py b/fasterai/_modidx.py index 47b25dbb..6e8d6471 100644 --- a/fasterai/_modidx.py +++ b/fasterai/_modidx.py @@ -341,8 +341,6 @@ 'fasterai/prune/prune_callback.py'), 'fasterai.prune.prune_callback.PruneCallback.__init__': ( 'prune/prune_callback.html#prunecallback.__init__', 'fasterai/prune/prune_callback.py'), - 'fasterai.prune.prune_callback.PruneCallback._build_pruning_schedule': ( 'prune/prune_callback.html#prunecallback._build_pruning_schedule', - 'fasterai/prune/prune_callback.py'), 'fasterai.prune.prune_callback.PruneCallback._rebind_opt': ( 'prune/prune_callback.html#prunecallback._rebind_opt', 'fasterai/prune/prune_callback.py'), 'fasterai.prune.prune_callback.PruneCallback._record_replacements': ( 'prune/prune_callback.html#prunecallback._record_replacements', @@ -358,8 +356,6 @@ 'fasterai.prune.pruner': { 'fasterai.prune.pruner.Pruner': ('prune/pruner.html#pruner', 'fasterai/prune/pruner.py'), 'fasterai.prune.pruner.Pruner.__init__': ( 'prune/pruner.html#pruner.__init__', 'fasterai/prune/pruner.py'), - 'fasterai.prune.pruner.Pruner._build_pruning_schedule': ( 'prune/pruner.html#pruner._build_pruning_schedule', - 'fasterai/prune/pruner.py'), 'fasterai.prune.pruner.Pruner._resolve_pruning_ratio_dict': ( 'prune/pruner.html#pruner._resolve_pruning_ratio_dict', 'fasterai/prune/pruner.py'), 'fasterai.prune.pruner.Pruner._to_tp_scheduler': ( 'prune/pruner.html#pruner._to_tp_scheduler', diff --git a/fasterai/prune/prune_callback.py b/fasterai/prune/prune_callback.py index c60aded8..cd64bd82 100644 --- a/fasterai/prune/prune_callback.py +++ b/fasterai/prune/prune_callback.py @@ -41,16 +41,6 @@ def __init__(self, self.extra_kwargs = kwargs self._validate_pruning_ratio() - def _build_pruning_schedule(self, sched_func): - "Create a schedule function compatible with torch-pruning's Pruner" - start_val, end_val = self.schedule.start_val, self.schedule.end_val - def scheduler(pruning_ratio, steps, start=start_val, end=end_val): - return [ - sched_func(start, end, i / float(steps)) * pruning_ratio - for i in range(steps + 1) - ] - return scheduler - def _validate_pruning_ratio(self) -> None: "Read pruning_ratio as a fraction, supporting both a single value and a per-layer dict" self._is_per_layer = isinstance(self.pruning_ratio, dict) @@ -71,20 +61,18 @@ def before_fit(self) -> None: self.example_inputs, _ = self.learn.dls.one_batch() - pruning_schedule = self._build_pruning_schedule(self.schedule.sched_func) - # nothing single to log for a per-layer dict: torch-pruning schedules each layer on its own - self.sparsity_levels = [] if self._is_per_layer else pruning_schedule(self.pruning_ratio, total_training_steps) - self.pruner = Pruner( self.learn.model, criteria=self.criteria, pruning_ratio=self.pruning_ratio, context=self.context, iterative_steps=total_training_steps, - schedule=pruning_schedule, + schedule=self.schedule, *self.extra_args, **self.extra_kwargs ) + # nothing single to log for a per-layer dict: torch-pruning schedules each layer on its own + self.sparsity_levels = [] if self._is_per_layer else list(self.pruner.pruner.per_step_pruning_ratio) def before_step(self) -> None: "Apply pruning before optimizer step" diff --git a/fasterai/prune/pruner.py b/fasterai/prune/pruner.py index 82858598..45101461 100644 --- a/fasterai/prune/pruner.py +++ b/fasterai/prune/pruner.py @@ -46,6 +46,7 @@ def regularize(self, model, alpha=2**4, bias=False): self._skipped = None # %% ../../nbs/prune/pruner.ipynb #63acddeb-f30e-448b-a397-d4cac2adba7a +from copy import copy from ..core.schedule import Schedule class Pruner(): @@ -106,20 +107,13 @@ def __init__(self, **kwargs ) - def _build_pruning_schedule(self, sched_func): - "Create a schedule function compatible with torch-pruning's Pruner" - def scheduler(pruning_ratio, steps, start=0, end=1): - return [ - sched_func(start, end, i / float(steps)) * pruning_ratio - for i in range(steps + 1) - ] - return scheduler - def _to_tp_scheduler(self, schedule): - "Convert Schedule object or callable to torch-pruning compatible scheduler" - if isinstance(schedule, Schedule): - return self._build_pruning_schedule(schedule.sched_func) - return schedule + "torch-pruning scheduler from a Schedule, `start_pct`/`end_pct` honoured; callables pass through" + if not isinstance(schedule, Schedule): return schedule + def scheduler(pruning_ratio, steps): + s = copy(schedule) + return [s.progress(i / steps) * pruning_ratio for i in range(steps + 1)] + return scheduler def _resolve_pruning_ratio_dict(self, ratio_dict): "Convert layer name strings to module references for torch-pruning" diff --git a/nbs/prune/prune_callback.ipynb b/nbs/prune/prune_callback.ipynb index cc318d8c..265b47aa 100644 --- a/nbs/prune/prune_callback.ipynb +++ b/nbs/prune/prune_callback.ipynb @@ -101,16 +101,6 @@ " self.extra_kwargs = kwargs\n", " self._validate_pruning_ratio()\n", "\n", - " def _build_pruning_schedule(self, sched_func):\n", - " \"Create a schedule function compatible with torch-pruning's Pruner\"\n", - " start_val, end_val = self.schedule.start_val, self.schedule.end_val\n", - " def scheduler(pruning_ratio, steps, start=start_val, end=end_val):\n", - " return [\n", - " sched_func(start, end, i / float(steps)) * pruning_ratio\n", - " for i in range(steps + 1)\n", - " ]\n", - " return scheduler\n", - "\n", " def _validate_pruning_ratio(self) -> None:\n", " \"Read pruning_ratio as a fraction, supporting both a single value and a per-layer dict\"\n", " self._is_per_layer = isinstance(self.pruning_ratio, dict)\n", @@ -131,20 +121,18 @@ "\n", " self.example_inputs, _ = self.learn.dls.one_batch()\n", "\n", - " pruning_schedule = self._build_pruning_schedule(self.schedule.sched_func)\n", - " # nothing single to log for a per-layer dict: torch-pruning schedules each layer on its own\n", - " self.sparsity_levels = [] if self._is_per_layer else pruning_schedule(self.pruning_ratio, total_training_steps)\n", - "\n", " self.pruner = Pruner(\n", " self.learn.model,\n", " criteria=self.criteria,\n", " pruning_ratio=self.pruning_ratio, \n", " context=self.context,\n", " iterative_steps=total_training_steps, \n", - " schedule=pruning_schedule,\n", + " schedule=self.schedule,\n", " *self.extra_args, \n", " **self.extra_kwargs\n", " )\n", + " # nothing single to log for a per-layer dict: torch-pruning schedules each layer on its own\n", + " self.sparsity_levels = [] if self._is_per_layer else list(self.pruner.pruner.per_step_pruning_ratio)\n", " \n", " def before_step(self) -> None:\n", " \"Apply pruning before optimizer step\"\n", @@ -237,7 +225,9 @@ " criteria=large_final\n", ")\n", "learn.fit(10, cbs=[cb])\n", - "```" + "```\n", + "\n", + "The schedule's `start_pct` and `end_pct` are honoured: `agp` prunes nothing before 20% of the fit, and `one_shot` prunes once, at the first step at or past half of it." ] }, { diff --git a/nbs/prune/pruner.ipynb b/nbs/prune/pruner.ipynb index 8428ebf3..8c922ddc 100644 --- a/nbs/prune/pruner.ipynb +++ b/nbs/prune/pruner.ipynb @@ -125,6 +125,7 @@ "outputs": [], "source": [ "#| export\n", + "from copy import copy\n", "from fasterai.core.schedule import Schedule\n", "\n", "class Pruner():\n", @@ -185,20 +186,13 @@ " **kwargs\n", " )\n", "\n", - " def _build_pruning_schedule(self, sched_func):\n", - " \"Create a schedule function compatible with torch-pruning's Pruner\"\n", - " def scheduler(pruning_ratio, steps, start=0, end=1):\n", - " return [\n", - " sched_func(start, end, i / float(steps)) * pruning_ratio\n", - " for i in range(steps + 1)\n", - " ]\n", - " return scheduler\n", - "\n", " def _to_tp_scheduler(self, schedule):\n", - " \"Convert Schedule object or callable to torch-pruning compatible scheduler\"\n", - " if isinstance(schedule, Schedule):\n", - " return self._build_pruning_schedule(schedule.sched_func)\n", - " return schedule\n", + " \"torch-pruning scheduler from a Schedule, `start_pct`/`end_pct` honoured; callables pass through\"\n", + " if not isinstance(schedule, Schedule): return schedule\n", + " def scheduler(pruning_ratio, steps):\n", + " s = copy(schedule)\n", + " return [s.progress(i / steps) * pruning_ratio for i in range(steps + 1)]\n", + " return scheduler\n", "\n", " def _resolve_pruning_ratio_dict(self, ratio_dict):\n", " \"Convert layer name strings to module references for torch-pruning\"\n", diff --git a/nbs/tests/test_prune_callback.ipynb b/nbs/tests/test_prune_callback.ipynb index 0331df0e..0f32d4e7 100644 --- a/nbs/tests/test_prune_callback.ipynb +++ b/nbs/tests/test_prune_callback.ipynb @@ -17,6 +17,7 @@ "source": [ "from fastcore.test import *\n", "from fastcore.foundation import L\n", + "from copy import copy\n", "from functools import partial\n", "from pathlib import Path\n", "import tempfile\n", @@ -138,7 +139,8 @@ "# (a) the optimizer holds every live parameter after every step — it held 1 of 10 before the rebind\n", "test_eq(_probe.missed, [0] * 6)\n", "# (b) the pruned conv keeps training after the prune — its weights did not move at all before the rebind\n", - "test_eq(_probe.channels, [8] * 6) # one_shot prunes at the first step and never again: same shape throughout\n", + "# one_shot (start_pct 0.5) prunes once, at the first step at or past half of the fit (the 3rd of 6)\n", + "test_eq(_probe.channels, [16, 16, 8, 8, 8, 8])\n", "assert _probe.snapshot is not None, \"the model was never pruned\"\n", "test_eq(_learn.model[0].weight.shape, _probe.snapshot.shape)\n", "assert not torch.equal(_learn.model[0].weight.detach(), _probe.snapshot), \\\n", @@ -169,6 +171,31 @@ " _learn.save('after_prune')" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "sched-levels", + "metadata": {}, + "outputs": [], + "source": [ + "# The callback's levels are the Schedule's, start_pct included\n", + "_cb = PruneCallback(pruning_ratio=0.5, schedule=agp, context='local', criteria=large_final)\n", + "_learn = _probe_learner(agp, _cb)\n", + "_learn.fit(3) # 9 steps\n", + "test_eq(_cb.sparsity_levels, [copy(agp).progress(i / 9) * 0.5 for i in range(10)])\n", + "test_eq(_cb.sparsity_levels, _cb.pruner.pruner.per_step_pruning_ratio)\n", + "test_eq(_cb.sparsity_levels[:2], [0., 0.])\n", + "\n", + "# A per-layer dict honours start_pct too and logs no single level\n", + "_probe = _OptProbe()\n", + "_cb = PruneCallback(pruning_ratio={'0': 0.5}, schedule=one_shot, context='local', criteria=large_final)\n", + "_learn = _probe_learner(one_shot, _cb)\n", + "_learn.add_cb(_probe)\n", + "_learn.fit(2)\n", + "test_eq(_probe.channels, [16, 16, 8, 8, 8, 8])\n", + "test_eq(_cb.sparsity_levels, [])" + ] + }, { "cell_type": "code", "execution_count": null, @@ -253,7 +280,7 @@ "outputs": [], "source": [ "# The state of a replaced parameter follows it, sliced exactly as the parameter was sliced. `iterative`\n", - "# prunes at steps 0, 3 and 6 of this fit, so the later prunes have a momentum to carry (the first has none)\n", + "# (start_pct 0.2) prunes at steps 1, 4 and 6 of this fit, each after an optimizer step, so each has a momentum to carry\n", "def _clone_state(s): return {k: v.clone() if torch.is_tensor(v) else v for k, v in s.items()}\n", "\n", "def _sliced(v, slices):\n", @@ -350,7 +377,7 @@ "for _n, _p in _learn.model.named_parameters():\n", " test_eq(_learn.opt.state[_p]['step'], _steps)\n", "\n", - "# and right after each of those prunes — every one but the first, which had no step behind it\n", + "# and right after each of those prunes — the first step, which had no step behind it, prunes nothing\n", "_recreated = [r['0.weight'][1] for r in _probe.records[1:] if '0.weight' in r]\n", "assert _recreated, \"agp pruned only at the first step\"\n", "for _s in _recreated:\n", diff --git a/nbs/tests/test_pruner.ipynb b/nbs/tests/test_pruner.ipynb index 1846c677..5c96b89b 100644 --- a/nbs/tests/test_pruner.ipynb +++ b/nbs/tests/test_pruner.ipynb @@ -21,6 +21,9 @@ "from torch.nn.utils import parametrize\n", "from fasterai.core.criteria import large_final\n", "from fasterai.core.parametrize import _unparametrize\n", + "from copy import copy\n", + "from torch_pruning.pruner.algorithms.scheduler import linear_scheduler\n", + "from fasterai.core.schedule import Schedule, agp, one_shot, one_cycle, cos, lin, sched_agp, sched_lin\n", "from fasterai.prune.pruner import *" ] }, @@ -348,6 +351,86 @@ " p.prune_model()\n", "test_eq(m(xb).shape, (2, 10))" ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "sched-levels", + "metadata": {}, + "outputs": [], + "source": [ + "# A Schedule reaches torch-pruning with its start_pct and end_pct\n", + "def _pruner(schedule, n=8, ratio=0.5, model=None, **kw):\n", + " return Pruner(_test_model() if model is None else model, ratio, 'local', large_final,\n", + " schedule=schedule, iterative_steps=n, example_inputs=x, **kw)\n", + "\n", + "def _levels(schedule, n=8, ratio=0.5, **kw): return _pruner(schedule, n, ratio, **kw).pruner.per_step_pruning_ratio\n", + "def _expected(s, n, r=0.5): return [copy(s).progress(i / n) * r for i in range(n + 1)]\n", + "\n", + "test_eq(_levels(one_shot), [0.] * 4 + [0.5] * 5) # the first step at or past start_pct 0.5\n", + "test_eq(_levels(agp, n=9), _expected(agp, 9))\n", + "test_eq(_levels(agp, n=9)[:2], [0., 0.])\n", + "_win = Schedule(sched_lin, start_pct=0.25, end_pct=0.75)\n", + "test_close([l / 0.5 for l in _levels(_win)], [0, 0, 0, .25, .5, .75, 1, 1, 1], eps=1e-9)\n", + "\n", + "# start_pct 0 and end_pct 1: the same levels as the formula used before these were read\n", + "for _s in (one_cycle, cos, lin, Schedule(sched_agp)):\n", + " test_eq(_levels(_s, n=9), [_s.sched_func(_s.start_val, _s.end_val, i / float(9)) * 0.5 for i in range(10)])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "sched-channels", + "metadata": {}, + "outputs": [], + "source": [ + "# The channels follow the levels: untouched before start_pct, constant after end_pct\n", + "def _channels(p, m, n=8):\n", + " \"`m[0]`'s channels after each of `n` prune steps of `p`\"\n", + " out = []\n", + " for _ in range(n):\n", + " p.prune_model()\n", + " out.append(m[0].out_channels)\n", + " return out\n", + "\n", + "_m = _test_model()\n", + "_ch = _channels(_pruner(_win, model=_m), _m)\n", + "test_eq(_ch[:2], [16, 16])\n", + "assert _ch[2] < 16, _ch\n", + "test_eq(len(set(_ch[5:])), 1)\n", + "test_eq(_ch[-1], 8)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "sched-reg-dict", + "metadata": {}, + "outputs": [], + "source": [ + "# The group-penalty pruner and a per-layer dict are scheduled the same way\n", + "test_eq(_levels(one_shot, reg=1e-4), _levels(one_shot))\n", + "\n", + "_m = _test_model()\n", + "_p = _pruner(one_shot, ratio={'0': 0.5}, model=_m)\n", + "test_eq(_p.pruner.pruning_ratio_dict[_m[0]], [0.] * 4 + [0.5] * 5)\n", + "test_eq(_channels(_p, _m), [16] * 3 + [8] * 5)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "sched-passthrough", + "metadata": {}, + "outputs": [], + "source": [ + "# A plain torch-pruning scheduler passes through, and the shared schedules are left untouched\n", + "test_eq(_levels(lambda r, n: [r] * (n + 1), n=3), [0.5] * 4)\n", + "test_eq(_levels(linear_scheduler), linear_scheduler(0.5, 8))\n", + "for _s in (one_shot, agp, lin):\n", + " test_eq((_s._current_progress, _s._previous_progress), (_s.start_val, _s.start_val))" + ] } ], "metadata": {