diff --git a/CHANGELOG.md b/CHANGELOG.md index 81124381..85a2ae71 100644 --- a/CHANGELOG.md +++ b/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` diff --git a/fasterai/prune/prune_callback.py b/fasterai/prune/prune_callback.py index e7157017..96f1657b 100644 --- a/fasterai/prune/prune_callback.py +++ b/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%}') diff --git a/nbs/prune/prune_callback.ipynb b/nbs/prune/prune_callback.ipynb index 38e6e8e0..d4c96ae3 100644 --- a/nbs/prune/prune_callback.ipynb +++ b/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%}')" ] }, diff --git a/nbs/tests/test_prune_callback.ipynb b/nbs/tests/test_prune_callback.ipynb index 2d79c8cd..de536a81 100644 --- a/nbs/tests/test_prune_callback.ipynb +++ b/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, diff --git a/nbs/tutorials/prune/prune_callback.ipynb b/nbs/tutorials/prune/prune_callback.ipynb index 2c236d69..08c9b56f 100644 --- a/nbs/tutorials/prune/prune_callback.ipynb +++ b/nbs/tutorials/prune/prune_callback.ipynb @@ -127,10 +127,10 @@ " \n", " \n", " 0\n", - " 0.601222\n", - " 0.380172\n", - " 0.845061\n", - " 00:05\n", + " 0.572080\n", + " 0.413692\n", + " 0.847091\n", + " 00:24\n", " \n", " \n", "" @@ -159,7 +159,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "after the warm-up epoch: 1249/1478 = 0.8451 Wilson 95% [0.8257, 0.8626]\n" + "after the warm-up epoch: 1252/1478 = 0.8471 Wilson 95% [0.8278, 0.8645]\n" ] } ], @@ -281,73 +281,73 @@ " \n", " \n", " 0\n", - " 0.354834\n", - " 0.316359\n", - " 0.871448\n", - " 00:06\n", + " 0.331890\n", + " 0.341828\n", + " 0.868742\n", + " 00:29\n", " \n", " \n", " 1\n", - " 0.289825\n", - " 0.298112\n", - " 0.886333\n", - " 00:06\n", + " 0.346365\n", + " 0.327003\n", + " 0.851150\n", + " 00:28\n", " \n", " \n", " 2\n", - " 0.245730\n", - " 0.237911\n", - " 0.900541\n", - " 00:07\n", + " 0.287846\n", + " 0.295884\n", + " 0.893775\n", + " 00:27\n", " \n", " \n", " 3\n", - " 0.228931\n", - " 0.332760\n", - " 0.872801\n", - " 00:07\n", + " 0.238572\n", + " 0.515673\n", + " 0.857916\n", + " 00:23\n", " \n", " \n", " 4\n", - " 0.184838\n", - " 0.218940\n", - " 0.914073\n", - " 00:08\n", + " 0.179014\n", + " 0.250834\n", + " 0.893099\n", + " 00:18\n", " \n", " \n", " 5\n", - " 0.213857\n", - " 0.227616\n", - " 0.915426\n", - " 00:07\n", + " 0.146992\n", + " 0.480372\n", + " 0.845737\n", + " 00:14\n", " \n", " \n", " 6\n", - " 0.266128\n", - " 0.235636\n", - " 0.897835\n", - " 00:08\n", + " 0.102759\n", + " 0.274859\n", + " 0.913396\n", + " 00:14\n", " \n", " \n", " 7\n", - " 0.241007\n", - " 0.234197\n", - " 0.904601\n", - " 00:06\n", + " 0.069884\n", + " 0.218804\n", + " 0.925575\n", + " 00:12\n", " \n", " \n", " 8\n", - " 0.230154\n", - " 0.237446\n", - " 0.900541\n", - " 00:05\n", + " 0.040234\n", + " 0.199090\n", + " 0.928281\n", + " 00:13\n", " \n", " \n", " 9\n", - " 0.226757\n", - " 0.228660\n", - " 0.909337\n", - " 00:05\n", + " 0.026111\n", + " 0.203525\n", + " 0.930311\n", + " 00:13\n", " \n", " \n", "" @@ -363,49 +363,49 @@ "name": "stdout", "output_type": "stream", "text": [ - "Pruning ratio at the end of epoch 0: 0.39%\n" + "Pruning ratio at the end of epoch 0: 0.40%\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ - "Pruning ratio at the end of epoch 1: 1.54%\n" + "Pruning ratio at the end of epoch 1: 1.57%\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ - "Pruning ratio at the end of epoch 2: 5.60%\n" + "Pruning ratio at the end of epoch 2: 5.68%\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ - "Pruning ratio at the end of epoch 3: 15.91%\n" + "Pruning ratio at the end of epoch 3: 16.06%\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ - "Pruning ratio at the end of epoch 4: 29.13%\n" + "Pruning ratio at the end of epoch 4: 29.25%\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ - "Pruning ratio at the end of epoch 5: 36.64%\n" + "Pruning ratio at the end of epoch 5: 36.69%\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ - "Pruning ratio at the end of epoch 6: 39.12%\n" + "Pruning ratio at the end of epoch 6: 39.14%\n" ] }, { @@ -469,22 +469,22 @@ "0.5.0.downsample.0 Conv2d 64 128 8,192 \n", "0.5.1.conv1 Conv2d 128 128 147,456 \n", "0.5.1.conv2 Conv2d 128 128 147,456 \n", - "0.6.0.conv1 Conv2d 128 239 275,328 \n", - "0.6.0.conv2 Conv2d 239 183 393,633 \n", - "0.6.0.downsample.0 Conv2d 128 183 23,424 \n", - "0.6.1.conv1 Conv2d 183 19 31,293 \n", - "0.6.1.conv2 Conv2d 19 183 31,293 \n", - "0.7.0.conv1 Conv2d 183 1 1,647 \n", + "0.6.0.conv1 Conv2d 128 224 258,048 \n", + "0.6.0.conv2 Conv2d 224 203 409,248 \n", + "0.6.0.downsample.0 Conv2d 128 203 25,984 \n", + "0.6.1.conv1 Conv2d 203 18 32,886 \n", + "0.6.1.conv2 Conv2d 18 203 32,886 \n", + "0.7.0.conv1 Conv2d 203 1 1,827 \n", "0.7.0.conv2 Conv2d 1 512 4,608 \n", - "0.7.0.downsample.0 Conv2d 183 512 93,696 \n", + "0.7.0.downsample.0 Conv2d 203 512 103,936 \n", "0.7.1.conv1 Conv2d 512 1 4,608 \n", "0.7.1.conv2 Conv2d 1 512 4,608 \n", - "1.4 Linear 1024 506 518,144 \n", - "1.8 Linear 506 2 1,012 \n", + "1.4 Linear 1024 502 514,048 \n", + "1.8 Linear 502 2 1,004 \n", "-------------------------------------------------------------------------------------\n", - "Total 2,064,446 \n", + "Total 2,074,843 \n", "Original 11,704,896 \n", - "Reduction 82.36%\n" + "Reduction 82.27%\n" ] } ], @@ -510,11 +510,11 @@ "name": "stdout", "output_type": "stream", "text": [ - "1.15 GMACs at 224x224, 2.07 M parameters, 3305 conv output channels\n", - "conv output channels removed: 31.1%\n", + "1.15 GMACs at 224x224, 2.08 M parameters, 3349 conv output channels\n", + "conv output channels removed: 30.2%\n", "parameter ratio: 0.18\n", - "MACs ratio at 224x224: 0.6328\n", - "MACs ratio at 64x64: 0.6339\n" + "MACs ratio at 224x224: 0.6335\n", + "MACs ratio at 64x64: 0.6346\n" ] } ], @@ -536,12 +536,12 @@ "metadata": {}, "source": [ "`pruning_ratio=0.4` asks the pruner for a global ratio of 0.4; the cells above show what that produced\n", - "here. Counting the output channels of every convolution, 4800 → 3305, 31.1% removed. In the report,\n", + "here. Counting the output channels of every convolution, 4800 → 3349, 30.2% removed. In the report,\n", "`0.4` to `0.7` are module paths and not ratios: the four ResNet stages inside the fastai body, two\n", "blocks each. The stem `0.0`, `0.4` and `0.5` keep their 64 and 128 channels, the removal falls on `0.6`\n", - "and `0.7` — down to 19 output channels in `0.6.1.conv1` and a single one in both `0.7.*.conv1` — and the\n", - "head `1.4` goes from 1024 to 506 units. Neither the parameter ratio (0.18) nor, in this run, the channel\n", - "ratio (31.1%) equals the 0.4 asked for: the ratio is what torch-pruning targets over its prunable groups,\n", + "and `0.7` — down to 18 output channels in `0.6.1.conv1` and a single one in both `0.7.*.conv1` — and the\n", + "head `1.4` goes from 512 to 502 units. Neither the parameter ratio (0.18) nor, in this run, the channel\n", + "ratio (30.2%) equals the 0.4 asked for: the ratio is what torch-pruning targets over its prunable groups,\n", "and a group cannot go below one channel — two convolutions here sit at that bound." ] }, @@ -555,7 +555,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "after pruning: 1344/1478 = 0.9093 Wilson 95% [0.8936, 0.9229]\n" + "after pruning: 1375/1478 = 0.9303 Wilson 95% [0.9162, 0.9422]\n" ] } ], @@ -573,7 +573,7 @@ "the 1478 validation images and printed by `report` with their Wilson 95% interval, but they come from\n", "different training budgets — one epoch against eleven — so they are not a dense-versus-pruned\n", "comparison. The absolute GMACs are for a 224x224 batch-1 input rather than the 64 px training\n", - "resolution; the MACs ratio barely moves with that choice, 0.6328 at 224x224 against 0.6339 at 64x64, as\n", + "resolution; the MACs ratio barely moves with that choice, 0.6335 at 224x224 against 0.6346 at 64x64, as\n", "the cell above prints. This page measures no latency." ] },