fix(prune)!: honour Schedule start_pct/end_pct in Pruner and PruneCallback

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
 

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',

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"

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"

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."
    ]
   },
   {

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",

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",

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": {