def test_multiple_pruning_callbacks(tmpdir, caplog, make_pruning_permanent: bool): model = TestModel() pruning_kwargs = { 'parameters_to_prune': [(model.layer.mlp_1, "weight"), (model.layer.mlp_3, "weight")], 'verbose': 2, "make_pruning_permanent": make_pruning_permanent } p1 = ModelPruning("l1_unstructured", amount=0.5, apply_pruning=lambda e: not e % 2, **pruning_kwargs) p2 = ModelPruning("random_unstructured", amount=0.25, apply_pruning=lambda e: e % 2, **pruning_kwargs) trainer = Trainer( default_root_dir=tmpdir, progress_bar_refresh_rate=0, weights_summary=None, checkpoint_callback=False, logger=False, limit_train_batches=10, limit_val_batches=2, max_epochs=3, callbacks=[p1, p2], ) with caplog.at_level(INFO): trainer.fit(model) actual = [m.strip() for m in caplog.messages] actual = [m for m in actual if m.startswith("Applied")] percentage = r"\(\d+(?:\.\d+)?%\)" expected = [ rf"Applied `L1Unstructured`. Pruned: \d+\/1122 {percentage} -> \d+\/1122 {percentage}", rf"Applied `L1Unstructured` to `Linear\(in_features=32, out_features=32, bias=True\).weight` with amount=0.5. Pruned: 0 \(0.00%\) -> \d+ {percentage}", # noqa: E501 rf"Applied `L1Unstructured` to `Linear\(in_features=32, out_features=2, bias=True\).weight` with amount=0.5. Pruned: 0 \(0.00%\) -> \d+ {percentage}", # noqa: E501 rf"Applied `RandomUnstructured`. Pruned: \d+\/1122 {percentage} -> \d+\/1122 {percentage}", rf"Applied `RandomUnstructured` to `Linear\(in_features=32, out_features=32, bias=True\).weight` with amount=0.25. Pruned: \d+ {percentage} -> \d+ {percentage}", # noqa: E501 rf"Applied `RandomUnstructured` to `Linear\(in_features=32, out_features=2, bias=True\).weight` with amount=0.25. Pruned: \d+ {percentage} -> \d+ {percentage}", # noqa: E501 rf"Applied `L1Unstructured`. Pruned: \d+\/1122 {percentage} -> \d+\/1122 {percentage}", rf"Applied `L1Unstructured` to `Linear\(in_features=32, out_features=32, bias=True\).weight` with amount=0.5. Pruned: \d+ {percentage} -> \d+ {percentage}", # noqa: E501 rf"Applied `L1Unstructured` to `Linear\(in_features=32, out_features=2, bias=True\).weight` with amount=0.5. Pruned: \d+ {percentage} -> \d+ {percentage}", # noqa: E501 ] expected = [re.compile(s) for s in expected] assert all(regex.match(s) for s, regex in zip(actual, expected)) filepath = str(tmpdir / "foo.ckpt") trainer.save_checkpoint(filepath) model.load_from_checkpoint(filepath, strict=False) has_pruning = hasattr(model.layer.mlp_1, "weight_orig") assert not has_pruning if make_pruning_permanent else has_pruning
def test_multiple_pruning_callbacks(tmpdir, caplog, make_pruning_permanent: bool): seed_everything(0) model = TestModel() pruning_kwargs = { 'parameters_to_prune': [(model.layer.mlp_1, "weight"), (model.layer.mlp_3, "weight")], 'verbose': 2, "make_pruning_permanent": make_pruning_permanent } p1 = ModelPruning("l1_unstructured", amount=0.5, apply_pruning=lambda e: not e % 2, **pruning_kwargs) p2 = ModelPruning("random_unstructured", amount=0.25, apply_pruning=lambda e: e % 2, **pruning_kwargs) trainer = Trainer( default_root_dir=tmpdir, progress_bar_refresh_rate=0, weights_summary=None, checkpoint_callback=False, logger=False, limit_train_batches=10, limit_val_batches=2, max_epochs=3, callbacks=[p1, p2], ) with caplog.at_level(INFO): trainer.fit(model) actual = [m.strip() for m in caplog.messages] actual = [m for m in actual if m.startswith("Applied")] assert actual == [ "Applied `L1Unstructured`. Pruned: 0/1122 (0.00%) -> 544/1122 (48.48%)", "Applied `L1Unstructured` to `Linear(in_features=32, out_features=32, bias=True).weight` with amount=0.5. Pruned: 0 (0.00%) -> 500 (48.83%)", # noqa: E501 "Applied `L1Unstructured` to `Linear(in_features=32, out_features=2, bias=True).weight` with amount=0.5. Pruned: 0 (0.00%) -> 44 (68.75%)", # noqa: E501 "Applied `RandomUnstructured`. Pruned: 544/1122 (48.48%) -> 680/1122 (60.61%)", "Applied `RandomUnstructured` to `Linear(in_features=32, out_features=32, bias=True).weight` with amount=0.25. Pruned: 500 (48.83%) -> 635 (62.01%)", # noqa: E501 "Applied `RandomUnstructured` to `Linear(in_features=32, out_features=2, bias=True).weight` with amount=0.25. Pruned: 44 (68.75%) -> 45 (70.31%)", # noqa: E501 "Applied `L1Unstructured`. Pruned: 680/1122 (60.61%) -> 884/1122 (78.79%)", "Applied `L1Unstructured` to `Linear(in_features=32, out_features=32, bias=True).weight` with amount=0.5. Pruned: 635 (62.01%) -> 830 (81.05%)", # noqa: E501 "Applied `L1Unstructured` to `Linear(in_features=32, out_features=2, bias=True).weight` with amount=0.5. Pruned: 45 (70.31%) -> 54 (84.38%)", # noqa: E501 ] filepath = str(tmpdir / "foo.ckpt") trainer.save_checkpoint(filepath) model.load_from_checkpoint(filepath, strict=False) has_pruning = hasattr(model.layer.mlp_1, "weight_orig") assert not has_pruning if make_pruning_permanent else has_pruning
def test_pruning_misconfiguration(): with pytest.raises(MisconfigurationException, match=r"chocolate isn't in \('weight', 'bias'\)"): ModelPruning(pruning_fn="l1_unstructured", parameter_names=["chocolate"]) with pytest.raises(MisconfigurationException, match=r"expected to be a str in \["): ModelPruning(pruning_fn={}) # noqa with pytest.raises(MisconfigurationException, match="should be provided"): ModelPruning(pruning_fn="random_structured") with pytest.raises(MisconfigurationException, match=r"must be any of \(0, 1, 2\)"): ModelPruning(pruning_fn="l1_unstructured", verbose=3) with pytest.raises(MisconfigurationException, match="requesting `ln_structured` pruning, the `pruning_norm`"): ModelPruning(pruning_fn="ln_structured", pruning_dim=0)
def train_with_pruning_callback( tmpdir, parameters_to_prune=False, use_global_unstructured=False, pruning_fn="l1_unstructured", use_lottery_ticket_hypothesis=False, accelerator=None, gpus=None, num_processes=1, ): model = TestModel() # Weights are random. None is 0 assert torch.all(model.layer.mlp_2.weight != 0) pruning_kwargs = { "pruning_fn": pruning_fn, "amount": 0.3, "use_global_unstructured": use_global_unstructured, "use_lottery_ticket_hypothesis": use_lottery_ticket_hypothesis, "verbose": 1, } if parameters_to_prune: pruning_kwargs["parameters_to_prune"] = [(model.layer.mlp_1, "weight"), (model.layer.mlp_2, "weight")] else: if isinstance(pruning_fn, str) and pruning_fn.endswith("_structured"): pruning_kwargs["parameter_names"] = ["weight"] else: pruning_kwargs["parameter_names"] = ["weight", "bias"] if isinstance(pruning_fn, str) and pruning_fn.endswith("_structured"): pruning_kwargs["pruning_dim"] = 0 if pruning_fn == "ln_structured": pruning_kwargs["pruning_norm"] = 1 # Misconfiguration checks if isinstance(pruning_fn, str) and pruning_fn.endswith( "_structured") and use_global_unstructured: with pytest.raises( MisconfigurationException, match="is supported with `use_global_unstructured=True`"): ModelPruning(**pruning_kwargs) return if ModelPruning._is_pruning_method( pruning_fn) and not use_global_unstructured: with pytest.raises(MisconfigurationException, match="currently only supported with"): ModelPruning(**pruning_kwargs) return pruning = ModelPruning(**pruning_kwargs) trainer = Trainer( default_root_dir=tmpdir, progress_bar_refresh_rate=0, weights_summary=None, checkpoint_callback=False, logger=False, limit_train_batches=10, limit_val_batches=2, max_epochs=10, accelerator=accelerator, gpus=gpus, num_processes=num_processes, callbacks=pruning, ) trainer.fit(model) trainer.test(model) if not accelerator: # Check some have been pruned assert torch.any(model.layer.mlp_2.weight == 0)