Merge branch 'master' into update/tpu

Author: kaushikb11

CHANGELOG.md

@@ -170,6 +170,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Sanitize `None` params during pruning ([#6836](https://github.com/PyTorchLightning/pytorch-lightning/pull/6836))
+
+
 - Made the `Plugin.reduce` method more consistent across all Plugins to reflect a mean-reduction by default ([#6011](https://github.com/PyTorchLightning/pytorch-lightning/pull/6011))
 
 
@@ -200,6 +203,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed TPU Colab hang issue, post training ([#6816](https://github.com/PyTorchLightning/pytorch-lightning/pull/6816))
 
 
+- Enforce an epoch scheduler interval when using SWA ([#6588](https://github.com/PyTorchLightning/pytorch-lightning/pull/6588))
+
+
 - Fixed an issue with `IterableDataset` when `__len__` is not defined ([#6828](https://github.com/PyTorchLightning/pytorch-lightning/pull/6828))
 
 

pytorch_lightning/callbacks/finetuning.py

@@ -77,7 +77,7 @@ def finetune_function(self, pl_module, current_epoch, optimizer, optimizer_idx):
                 # When `current_epoch` is 10, feature_extractor will start training.
                 if current_epoch == self._unfreeze_at_epoch:
                     self.unfreeze_and_add_param_group(
-                        module=pl_module.feature_extractor,
+                        modules=pl_module.feature_extractor,
                         optimizer=optimizer,
                         train_bn=True,
                     )

pytorch_lightning/callbacks/pruning.py

@@ -422,7 +422,9 @@ def sanitize_parameters_to_prune(
         current_modules = [m for m in pl_module.modules() if not isinstance(m, _MODULE_CONTAINERS)]
 
         if parameters_to_prune is None:
-            parameters_to_prune = [(m, p) for p in parameters for m in current_modules if hasattr(m, p)]
+            parameters_to_prune = [
+                (m, p) for p in parameters for m in current_modules if getattr(m, p, None) is not None
+            ]
         elif (
             isinstance(parameters_to_prune, (list, tuple)) and len(parameters_to_prune) > 0
             and all(len(p) == 2 for p in parameters_to_prune)

pytorch_lightning/callbacks/stochastic_weight_avg.py

@@ -187,14 +187,15 @@ def on_train_epoch_start(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningMo
                 anneal_strategy=self._annealing_strategy,
                 last_epoch=trainer.max_epochs if self._annealing_strategy == "cos" else -1
             )
+            _scheduler_config = _get_default_scheduler_config()
+            assert _scheduler_config["interval"] == "epoch" and _scheduler_config["frequency"] == 1
+            _scheduler_config["scheduler"] = self._swa_scheduler
 
             if trainer.lr_schedulers:
                 lr_scheduler = trainer.lr_schedulers[0]["scheduler"]
                 rank_zero_warn(f"Swapping lr_scheduler {lr_scheduler} for {self._swa_scheduler}")
-                trainer.lr_schedulers[0]["scheduler"] = self._swa_scheduler
+                trainer.lr_schedulers[0] = _scheduler_config
             else:
-                _scheduler_config = _get_default_scheduler_config()
-                _scheduler_config["scheduler"] = self._swa_scheduler
                 trainer.lr_schedulers.append(_scheduler_config)
 
             self.n_averaged = torch.tensor(0, dtype=torch.long, device=pl_module.device)

tests/callbacks/test_pruning.py

@@ -36,7 +36,7 @@ def __init__(self):
         self.layer = Sequential(
             OrderedDict([
                 ("mlp_1", nn.Linear(32, 32)),
-                ("mlp_2", nn.Linear(32, 32)),
+                ("mlp_2", nn.Linear(32, 32, bias=False)),
                 ("mlp_3", nn.Linear(32, 2)),
             ])
         )
@@ -85,7 +85,10 @@ def train_with_pruning_callback(
     if parameters_to_prune:
         pruning_kwargs["parameters_to_prune"] = [(model.layer.mlp_1, "weight"), (model.layer.mlp_2, "weight")]
     else:
-        pruning_kwargs["parameter_names"] = ["weight"]
+        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":
@@ -249,14 +252,14 @@ def test_multiple_pruning_callbacks(tmpdir, caplog, make_pruning_permanent: bool
     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%) -> 506 (49.41%)",  # noqa: E501
-        "Applied `L1Unstructured` to `Linear(in_features=32, out_features=2, bias=True).weight` with amount=0.5. Pruned: 0 (0.00%) -> 38 (59.38%)",  # noqa: E501
+        "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: 506 (49.41%) -> 633 (61.82%)",  # noqa: E501
-        "Applied `RandomUnstructured` to `Linear(in_features=32, out_features=2, bias=True).weight` with amount=0.25. Pruned: 38 (59.38%) -> 47 (73.44%)",  # noqa: E501
+        "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: 633 (61.82%) -> 828 (80.86%)",  # noqa: E501
-        "Applied `L1Unstructured` to `Linear(in_features=32, out_features=2, bias=True).weight` with amount=0.5. Pruned: 47 (73.44%) -> 56 (87.50%)",  # noqa: E501
+        "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")

tests/callbacks/test_stochastic_weight_avg.py

@@ -27,16 +27,18 @@
 
 if _TORCH_GREATER_EQUAL_1_6:
     from pytorch_lightning.callbacks import StochasticWeightAveraging
+    from torch.optim.swa_utils import SWALR
 
     class SwaTestModel(BoringModel):
 
-        def __init__(self, batchnorm: bool = True):
+        def __init__(self, batchnorm: bool = True, interval: str = "epoch"):
             super().__init__()
             layers = [nn.Linear(32, 32)]
             if batchnorm:
                 layers.append(nn.BatchNorm1d(32))
             layers += [nn.ReLU(), nn.Linear(32, 2)]
             self.layer = nn.Sequential(*layers)
+            self.interval = interval
 
         def training_step(self, batch, batch_idx):
             output = self.forward(batch)
@@ -46,6 +48,14 @@ def training_step(self, batch, batch_idx):
         def train_dataloader(self):
             return DataLoader(RandomDataset(32, 64), batch_size=2)
 
+        def configure_optimizers(self):
+            optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
+            return {
+                "optimizer": optimizer,
+                "scheduler": torch.optim.lr_scheduler.StepLR(optimizer, step_size=1),
+                "interval": self.interval,
+            }
+
     class SwaTestCallback(StochasticWeightAveraging):
         update_parameters_calls: int = 0
         transfer_weights_calls: int = 0
@@ -61,6 +71,10 @@ def transfer_weights(self, *args, **kwargs):
         def on_train_epoch_start(self, trainer, *args):
             super().on_train_epoch_start(trainer, *args)
             assert trainer.train_loop._skip_backward == (trainer.current_epoch > self.swa_end)
+            if self.swa_start <= trainer.current_epoch:
+                assert isinstance(trainer.lr_schedulers[0]["scheduler"], SWALR)
+                assert trainer.lr_schedulers[0]["interval"] == "epoch"
+                assert trainer.lr_schedulers[0]["frequency"] == 1
 
         def on_train_epoch_end(self, trainer, *args):
             super().on_train_epoch_end(trainer, *args)
@@ -89,8 +103,8 @@ def on_train_end(self, trainer, pl_module):
 
 
 @mock.patch.dict(os.environ, {"PL_DEV_DEBUG": "1"})
-def train_with_swa(tmpdir, batchnorm=True, accelerator=None, gpus=None, num_processes=1):
-    model = SwaTestModel(batchnorm=batchnorm)
+def train_with_swa(tmpdir, batchnorm=True, accelerator=None, gpus=None, num_processes=1, interval="epoch"):
+    model = SwaTestModel(batchnorm=batchnorm, interval=interval)
     swa_start = 2
     max_epochs = 5
     swa_callback = SwaTestCallback(swa_epoch_start=swa_start, swa_lrs=0.1)
@@ -140,6 +154,12 @@ def test_swa_callback(tmpdir, batchnorm: bool):
     train_with_swa(tmpdir, batchnorm=batchnorm)
 
 
+@RunIf(min_torch="1.6.0")
+@pytest.mark.parametrize("interval", ("epoch", "step"))
+def test_swa_callback_scheduler_step(tmpdir, interval: bool):
+    train_with_swa(tmpdir, interval=interval)
+
+
 @RunIf(min_torch="1.6.0")
 def test_swa_raises():
     with pytest.raises(MisconfigurationException, match=">0 integer or a float between 0 and 1"):