add new test

Author: tchaton

pytorch_lightning/trainer/trainer.py

@@ -459,6 +459,7 @@ def fit(
         # SET UP TRAINING
         # ----------------------------
         self.call_setup_hook(model)
+        self.call_hook("on_before_accelerator_backend_setup", model)
         self.accelerator_backend.setup(self, model)
         self.setup_trainer(model)
 
@@ -470,7 +471,6 @@ def fit(
 
         # plugin will setup training (e.g. ddp will launch child processes)
         # TODO: the old setup is now called "pre_training", where should this hook be called now?
-        self.call_hook("on_before_accelerator_backend_setup", model)
         self.training_type_plugin.pre_training()
         self.precision_plugin.pre_training()
 

tests/callbacks/test_finetuning_callback.py

@@ -19,6 +19,7 @@
 
 from pytorch_lightning import LightningModule, seed_everything, Trainer
 from pytorch_lightning.callbacks import BackboneFinetuning, BaseFinetuning
+from pytorch_lightning.callbacks.base import Callback
 from tests.base import BoringModel, RandomDataset
 
 
@@ -215,3 +216,31 @@ def __init__(self):
             assert torch.equal(optimizer.param_groups[2]["params"][0], model.backbone[2].weight)
             assert torch.equal(optimizer.param_groups[2]["params"][1], model.backbone[3].weight)
             assert torch.equal(optimizer.param_groups[2]["params"][2], model.backbone[4].weight)
+
+
+def test_on_before_accelerator_backend_setup(tmpdir):
+    """
+    `on_before_accelerator_backend_setup` hook is used make sure the finetuning freeze call is made
+    before configure_optimizers call.
+    """
+
+    class TestCallback(Callback):
+
+        def on_before_accelerator_backend_setup(self, trainer, pl_module):
+            pl_module.on_before_accelerator_backend_setup_called = True
+
+    class TestModel(BoringModel):
+
+        def __init__(self):
+            super().__init__()
+            self.on_before_accelerator_backend_setup_called = False
+
+        def configure_optimizers(self):
+            assert self.on_before_accelerator_backend_setup_called
+            return super().configure_optimizers()
+
+    model = TestModel()
+    callback = TestCallback()
+
+    trainer = Trainer(default_root_dir=tmpdir, callbacks=[callback], fast_dev_run=True)
+    trainer.fit(model)