fix(prune): PruneCallback logs the applied ratio and traces on the fit's own batch

CHANGELOG.md

@@ -14,6 +14,8 @@
 
 ### 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
+- `PruneCallback` logs, at the end of each epoch, the pruning ratio torch-pruning actually applied. torch-pruning's `step()` advances its step counter before reading the level, so after `k` steps the model is at level `k`; the log printed level `k - 1`, one step behind, and the last epoch did not print the final target
+- `PruneCallback` hands the batch it takes from `learn.dls.one_batch()` to the `Pruner` as `example_inputs`, so the dependency graph is traced on the fit's own data. It was read and then dropped, and the `Pruner` traced on its default `torch.randn(1, 3, 224, 224)`, which failed on a model that does not accept a 224x224 image. An `example_inputs` passed to the callback still wins
 
 ### Dependencies
 - `torch-pruning>=1.5`

fasterai/prune/prune_callback.py

@@ -69,8 +69,8 @@ def before_fit(self) -> None:
             context=self.context,
             iterative_steps=total_training_steps, 
             schedule=self.schedule,
-            *self.extra_args, 
-            **self.extra_kwargs
+            *self.extra_args,
+            **{'example_inputs': self.example_inputs, **self.extra_kwargs}  # an example_inputs given by the user wins
         )
         # 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)
@@ -122,6 +122,6 @@ def after_epoch(self) -> None:
             print(f'Pruning {len(self.pruning_ratio)} layers to per-layer targets (epoch {self.epoch})')
             return
         completed_steps = (self.epoch + 1) * len(self.learn.dls.train)
-        if completed_steps > 0 and completed_steps <= len(self.sparsity_levels):
-            current_ratio = self.sparsity_levels[completed_steps - 1]
+        if completed_steps < len(self.sparsity_levels):  # torch-pruning's k-th step applies levels[k]
+            current_ratio = self.sparsity_levels[completed_steps]
             print(f'Pruning ratio at the end of epoch {self.epoch}: {current_ratio:.2%}')

nbs/prune/prune_callback.ipynb

@@ -129,8 +129,8 @@
     "            context=self.context,\n",
     "            iterative_steps=total_training_steps, \n",
     "            schedule=self.schedule,\n",
-    "            *self.extra_args, \n",
-    "            **self.extra_kwargs\n",
+    "            *self.extra_args,\n",
+    "            **{'example_inputs': self.example_inputs, **self.extra_kwargs}  # an example_inputs given by the user wins\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",
@@ -182,8 +182,8 @@
     "            print(f'Pruning {len(self.pruning_ratio)} layers to per-layer targets (epoch {self.epoch})')\n",
     "            return\n",
     "        completed_steps = (self.epoch + 1) * len(self.learn.dls.train)\n",
-    "        if completed_steps > 0 and completed_steps <= len(self.sparsity_levels):\n",
-    "            current_ratio = self.sparsity_levels[completed_steps - 1]\n",
+    "        if completed_steps < len(self.sparsity_levels):  # torch-pruning's k-th step applies levels[k]\n",
+    "            current_ratio = self.sparsity_levels[completed_steps]\n",
     "            print(f'Pruning ratio at the end of epoch {self.epoch}: {current_ratio:.2%}')"
    ]
   },

nbs/tests/test_prune_callback.ipynb

@@ -30,7 +30,9 @@
     "from fastai.optimizer import OptimWrapper, SGD\n",
     "from fasterai.core.criteria import large_final\n",
     "from fasterai.core.schedule import agp, iterative, one_shot\n",
-    "from fasterai.prune.prune_callback import *"
+    "from fasterai.prune.prune_callback import *\n",
+    "from contextlib import redirect_stdout\n",
+    "import io\n"
    ]
   },
   {
@@ -208,6 +210,60 @@
     "test_eq(_cb.sparsity_levels, [])"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "epoch-log",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Each epoch logs the level torch-pruning applied at its last step: `step()` advances its counter before\n",
+    "# reading the level, so after `k` steps the model is at `levels[k]` and the last epoch logs the target itself\n",
+    "class _LevelProbe(Callback):\n",
+    "    \"Record, at the end of each epoch, the level torch-pruning last applied\"\n",
+    "    order = 60\n",
+    "    def __init__(self, cb): self.cb, self.applied = cb, []\n",
+    "    def after_epoch(self):\n",
+    "        tp = self.cb.pruner.pruner\n",
+    "        self.applied.append(tp.per_step_pruning_ratio[tp.current_step])\n",
+    "\n",
+    "_cb = PruneCallback(pruning_ratio=0.5, schedule=agp, context='local', criteria=large_final)\n",
+    "_learn = _probe_learner(agp, _cb)\n",
+    "_lp = _LevelProbe(_cb)\n",
+    "_learn.add_cb(_lp)\n",
+    "with redirect_stdout(io.StringIO()) as _out: _learn.fit(3)\n",
+    "_logged = [l for l in _out.getvalue().splitlines() if l.startswith('Pruning ratio')]\n",
+    "test_eq(_logged, [f'Pruning ratio at the end of epoch {e}: {r:.2%}' for e, r in enumerate(_lp.applied)])\n",
+    "test_eq(_logged[-1], 'Pruning ratio at the end of epoch 2: 50.00%')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "trace-inputs",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# The pruner traces the model on a batch of the fit's own data, not on a 224x224 image: this head only\n",
+    "# accepts 8x8 inputs, and an `example_inputs` given by the user wins over the batch\n",
+    "def _flat_head():\n",
+    "    torch.manual_seed(0)\n",
+    "    return nn.Sequential(nn.Conv2d(3, 16, 3, padding=1), nn.BatchNorm2d(16), nn.ReLU(),\n",
+    "                         nn.Conv2d(16, 16, 3, padding=1), nn.ReLU(), nn.Flatten(), nn.Linear(16 * 8 * 8, 10))\n",
+    "\n",
+    "_cb = PruneCallback(pruning_ratio=0.5, schedule=one_shot, context='local', criteria=large_final)\n",
+    "_learn = _probe_learner(one_shot, _cb, model=_flat_head())\n",
+    "_learn.fit(1)\n",
+    "test_eq(tuple(_cb.pruner.example_inputs.shape), (16, 3, 8, 8))\n",
+    "test_eq(_cb.pruner.example_inputs.device, next(_learn.model.parameters()).device)\n",
+    "assert _learn.model[0].out_channels < 16, \"the model was never pruned\"\n",
+    "\n",
+    "_cb = PruneCallback(pruning_ratio=0.5, schedule=one_shot, context='local', criteria=large_final,\n",
+    "                    example_inputs=torch.randn(2, 3, 8, 8))\n",
+    "_probe_learner(one_shot, _cb, model=_flat_head()).fit(1)\n",
+    "test_eq(tuple(_cb.pruner.example_inputs.shape), (2, 3, 8, 8))\n"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,