fix restoring finetune callbacks after accelerator setup on training resume (#8501)

Co-authored-by: tchaton <thomas@grid.ai>
Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 5 people

CHANGELOG.md

@@ -547,6 +547,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed a `TypeError` when wrapping optimizers in the `HorovodPlugin` and running `Trainer.test` ([#7840](https://github.com/PyTorchLightning/pytorch-lightning/pull/7840))
 
 
+- Fixed `BackboneFinetuning` restoration ([#8501](https://github.com/PyTorchLightning/pytorch-lightning/pull/8501))
+
+
+
 ## [1.3.8] - 2021-07-01
 
 ### Fixed

pytorch_lightning/callbacks/finetuning.py

@@ -82,25 +82,33 @@ def finetune_function(self, pl_module, current_epoch, optimizer, optimizer_idx):
     """
 
     def __init__(self):
-        self._internal_state: Dict[int, List[Dict[str, Any]]] = {}
+        self._internal_optimizer_metadata: Dict[int, List[Dict[str, Any]]] = {}
+        self._restarting = False
 
     def on_save_checkpoint(
         self,
         trainer: 'pl.Trainer',
         pl_module: 'pl.LightningModule',
         checkpoint: Dict[str, Any],
     ) -> Dict[int, List[Dict[str, Any]]]:
-        return self._internal_state
+        return self._internal_optimizer_metadata
 
     def on_load_checkpoint(
         self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', callback_state: Dict[int, List[Dict[str, Any]]]
     ) -> None:
-        self._internal_state = callback_state
+        self._restarting = True
+        self._internal_optimizer_metadata = callback_state
+
+    def on_fit_start(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         # restore the param_groups created during the previous training.
-        named_parameters = dict(pl_module.named_parameters())
-        for opt_idx, optimizer in enumerate(trainer.optimizers):
-            param_groups = self.__apply_mapping_to_param_groups(self._internal_state[opt_idx], named_parameters)
-            optimizer.param_groups = param_groups
+        if self._restarting:
+            named_parameters = dict(pl_module.named_parameters())
+            for opt_idx, optimizer in enumerate(trainer.optimizers):
+                param_groups = self.__apply_mapping_to_param_groups(
+                    self._internal_optimizer_metadata[opt_idx], named_parameters
+                )
+                optimizer.param_groups = param_groups
+            self._restarting = False
 
     @staticmethod
     def flatten_modules(modules: Union[Module, Iterable[Union[Module, Iterable]]]) -> List[Module]:
@@ -278,11 +286,13 @@ def _store(
         current_param_groups: List[Dict[str, Any]],
     ) -> None:
         mapping = {p: n for n, p in pl_module.named_parameters()}
-        if opt_idx not in self._internal_state:
-            self._internal_state[opt_idx] = self.__apply_mapping_to_param_groups(current_param_groups, mapping)
+        if opt_idx not in self._internal_optimizer_metadata:
+            self._internal_optimizer_metadata[opt_idx] = self.__apply_mapping_to_param_groups(
+                current_param_groups, mapping
+            )
         elif num_param_groups != len(current_param_groups):
             # save new param_groups possibly created by the users.
-            self._internal_state[opt_idx].extend(
+            self._internal_optimizer_metadata[opt_idx].extend(
                 self.__apply_mapping_to_param_groups(current_param_groups[num_param_groups:], mapping)
             )
 
@@ -362,15 +372,33 @@ def __init__(
     ):
         super().__init__()
 
-        self.unfreeze_backbone_at_epoch = unfreeze_backbone_at_epoch
-        self.backbone_initial_lr = backbone_initial_lr
-        self.lambda_func = lambda_func
-        self.backbone_initial_ratio_lr = backbone_initial_ratio_lr
-        self.should_align = should_align
-        self.initial_denom_lr = initial_denom_lr
-        self.train_bn = train_bn
-        self.round = round
-        self.verbose = verbose
+        self.unfreeze_backbone_at_epoch: int = unfreeze_backbone_at_epoch
+        self.lambda_func: Callable = lambda_func
+        self.backbone_initial_ratio_lr: float = backbone_initial_ratio_lr
+        self.backbone_initial_lr: Optional[float] = backbone_initial_lr
+        self.should_align: bool = should_align
+        self.initial_denom_lr: float = initial_denom_lr
+        self.train_bn: bool = train_bn
+        self.verbose: bool = verbose
+        self.round: int = round
+        self.previous_backbone_lr: Optional[float] = None
+
+    def on_save_checkpoint(
+        self,
+        trainer: 'pl.Trainer',
+        pl_module: 'pl.LightningModule',
+        checkpoint: Dict[str, Any],
+    ) -> Dict[int, Any]:
+        return {
+            "internal_optimizer_metadata": self._internal_optimizer_metadata,
+            "previous_backbone_lr": self.previous_backbone_lr
+        }
+
+    def on_load_checkpoint(
+        self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', callback_state: Dict[int, List[Dict[str, Any]]]
+    ) -> None:
+        self.previous_backbone_lr = callback_state["previous_backbone_lr"]
+        super().on_load_checkpoint(trainer, pl_module, callback_state["internal_optimizer_metadata"])
 
     def on_fit_start(self, trainer, pl_module):
         """
@@ -379,7 +407,7 @@ def on_fit_start(self, trainer, pl_module):
                 If LightningModule has no nn.Module `backbone` attribute.
         """
         if hasattr(pl_module, "backbone") and isinstance(pl_module.backbone, Module):
-            return
+            return super().on_fit_start(trainer, pl_module)
         raise MisconfigurationException("The LightningModule should have a nn.Module `backbone` attribute")
 
     def freeze_before_training(self, pl_module: 'pl.LightningModule'):

tests/callbacks/test_finetuning_callback.py

@@ -236,7 +236,7 @@ def finetune_function(self, pl_module: LightningModule, epoch: int, optimizer: O
         self.unfreeze_and_add_param_group(pl_module.layer[epoch + 1], optimizer)
 
 
-def test_base_finetuning_internal_state(tmpdir):
+def test_base_finetuning_internal_optimizer_metadata(tmpdir):
     """Test the param_groups updates are properly saved within the internal state of the BaseFinetuning Callbacks"""
 
     seed_everything(42)
@@ -265,18 +265,18 @@ def configure_optimizers(self):
     model = FreezeModel()
     trainer = Trainer(default_root_dir=tmpdir, max_epochs=5, limit_train_batches=1, callbacks=[cb, chk])
     trainer.fit(model)
-    assert len(cb._internal_state[0]) == 6
-    assert cb._internal_state[0][0]["params"] == ['layer.0.weight']
-    assert cb._internal_state[0][1]["params"] == ['layer.1.weight', 'layer.1.bias']
-    assert cb._internal_state[0][2]["params"] == ['layer.2.weight']
-    assert cb._internal_state[0][3]["params"] == ['layer.3.weight', 'layer.3.bias']
-    assert cb._internal_state[0][4]["params"] == ['layer.4.weight']
-    assert cb._internal_state[0][5]["params"] == ['layer.5.weight', 'layer.5.bias']
+    assert len(cb._internal_optimizer_metadata[0]) == 6
+    assert cb._internal_optimizer_metadata[0][0]["params"] == ['layer.0.weight']
+    assert cb._internal_optimizer_metadata[0][1]["params"] == ['layer.1.weight', 'layer.1.bias']
+    assert cb._internal_optimizer_metadata[0][2]["params"] == ['layer.2.weight']
+    assert cb._internal_optimizer_metadata[0][3]["params"] == ['layer.3.weight', 'layer.3.bias']
+    assert cb._internal_optimizer_metadata[0][4]["params"] == ['layer.4.weight']
+    assert cb._internal_optimizer_metadata[0][5]["params"] == ['layer.5.weight', 'layer.5.bias']
 
     model = FreezeModel()
     cb = OnEpochLayerFinetuning()
     trainer = Trainer(max_epochs=10, resume_from_checkpoint=chk.last_model_path, callbacks=[cb])
-    with pytest.raises(ValueError, match="loaded state dict has a different number of parameter groups"):
+    with pytest.raises(IndexError, match="index 6 is out of range"):
         trainer.fit(model)
 
 
@@ -365,3 +365,115 @@ def forward(self, x):
     # conv0.weight, conv0.bias, bn0.weight, bn0.bias, parent_param
     # conv1.weight, conv1.bias, bn1.weight, bn1.bias
     assert len(encoder_params) == 9
+
+
+class TestCallbacksRestoreCallback(BaseFinetuning):
+
+    def freeze_before_training(self, pl_module):
+        self.freeze(pl_module.layer[:3])
+
+    def finetune_function(self, pl_module, epoch, optimizer, opt_idx):
+        if epoch >= 1:
+            self.unfreeze_and_add_param_group(pl_module.layer[epoch - 1], optimizer)
+
+
+class FinetuningBoringModel(BoringModel):
+
+    def __init__(self):
+        super().__init__()
+        self.layer = nn.Sequential(nn.Linear(32, 32), nn.Linear(32, 32), nn.Linear(32, 32), nn.Linear(32, 2))
+
+    def configure_optimizers(self):
+        parameters = filter(lambda x: x.requires_grad, self.parameters())
+        optimizer = torch.optim.SGD(parameters, lr=0.1)
+        return optimizer
+
+
+def test_callbacks_restore(tmpdir):
+    """
+    Test callbacks restore is called after optimizers have been re-created
+    but before optimizer states reload
+    """
+    chk = ModelCheckpoint(dirpath=tmpdir, save_last=True)
+
+    model = FinetuningBoringModel()
+    callback = TestCallbacksRestoreCallback()
+
+    trainer_kwargs = dict(
+        default_root_dir=tmpdir, limit_train_batches=1, limit_val_batches=1, callbacks=[callback, chk], max_epochs=2
+    )
+
+    trainer = Trainer(**trainer_kwargs)
+    trainer.fit(model)
+
+    # only 1 optimizer
+    assert len(callback._internal_optimizer_metadata) == 1
+
+    # only 2 param groups
+    assert len(callback._internal_optimizer_metadata[0]) == 2
+
+    # original parameters
+    assert callback._internal_optimizer_metadata[0][0] == {
+        'lr': 0.1,
+        'momentum': 0,
+        'dampening': 0,
+        'weight_decay': 0,
+        'nesterov': False,
+        'params': ['layer.3.weight', 'layer.3.bias']
+    }
+
+    # new param group
+    assert callback._internal_optimizer_metadata[0][1] == {
+        'lr': 0.01,
+        'momentum': 0,
+        'dampening': 0,
+        'weight_decay': 0,
+        'nesterov': False,
+        'params': ['layer.0.weight', 'layer.0.bias']
+    }
+
+    trainer_kwargs["max_epochs"] = 3
+    trainer_kwargs["resume_from_checkpoint"] = chk.last_model_path
+
+    trainer = Trainer(**trainer_kwargs)
+    trainer.fit(model)
+
+
+def test_callbacks_restore_backbone(tmpdir):
+    """
+    Test callbacks restore is called after optimizers have been re-created
+    but before optimizer states reload
+    """
+
+    class BackboneBoringModel(BoringModel):
+
+        def __init__(self):
+            super().__init__()
+            self.layer = nn.Linear(32, 2)
+            self.backbone = nn.Linear(32, 32)
+
+        def forward(self, x):
+            return self.layer(self.backbone(x))
+
+    ckpt = ModelCheckpoint(dirpath=tmpdir, save_last=True)
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        limit_train_batches=1,
+        limit_val_batches=1,
+        max_epochs=2,
+        progress_bar_refresh_rate=0,
+        callbacks=[ckpt, BackboneFinetuning(unfreeze_backbone_at_epoch=1)]
+    )
+    trainer.fit(BackboneBoringModel())
+
+    # initialize a trainer that continues the previous training
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        limit_train_batches=1,
+        limit_val_batches=1,
+        max_epochs=3,
+        progress_bar_refresh_rate=0,
+        callbacks=BackboneFinetuning(unfreeze_backbone_at_epoch=1),
+        resume_from_checkpoint=ckpt.last_model_path
+    )
+    trainer.fit(BackboneBoringModel())

tests/checkpointing/test_trainer_checkpoint.py

@@ -91,15 +91,15 @@ def test_accumulated_gradient_batches_with_resume_from_checkpoint(tmpdir):
     This test validates that accumulated gradient is properly recomputed and reset on the trainer.
     """
 
-    cb = ModelCheckpoint(dirpath=tmpdir, save_last=True)
+    ckpt = ModelCheckpoint(dirpath=tmpdir, save_last=True)
     model = BoringModel()
     trainer_kwargs = dict(
-        max_epochs=1, accumulate_grad_batches={0: 2}, callbacks=cb, limit_train_batches=1, limit_val_batches=0
+        max_epochs=1, accumulate_grad_batches={0: 2}, callbacks=ckpt, limit_train_batches=1, limit_val_batches=0
     )
     trainer = Trainer(**trainer_kwargs)
     trainer.fit(model)
 
     trainer_kwargs['max_epochs'] = 2
-    trainer_kwargs['resume_from_checkpoint'] = cb.last_model_path
+    trainer_kwargs['resume_from_checkpoint'] = ckpt.last_model_path
     trainer = Trainer(**trainer_kwargs)
     trainer.fit(model)