refactor finetuningscheduler to remove core pl trainer convenience attribute and EarlyStopping modification. FTSEarlyStopping extension of EarlyStopping added. FTS dependencies managed by a callbackresolvermixin

Author: speediedan

pl_examples/basic_examples/config/fts/fts_explicit.yaml

@@ -1,14 +1,14 @@
 trainer:
   callbacks:
-  - class_path: pytorch_lightning.callbacks.FinetuningScheduler
+  - class_path: pytorch_lightning.callbacks.finetuning_scheduler.FinetuningScheduler
     init_args:
       ft_schedule: ./pl_examples/basic_examples/config/fts/RteBoolqModule_ft_schedule_albert_base.yaml
   - class_path: pytorch_lightning.callbacks.finetuning_scheduler.FTSCheckpoint
     init_args:
       save_top_k: 5
       monitor: val_loss
       verbose: true
-  - class_path: pytorch_lightning.callbacks.EarlyStopping
+  - class_path: pytorch_lightning.callbacks.finetuning_scheduler.FTSEarlyStopping
     init_args:
       monitor: val_loss
       min_delta: 0.001

pl_examples/basic_examples/config/fts/fts_implicit.yaml

@@ -1,12 +1,12 @@
 trainer:
   callbacks:
-  - class_path: pytorch_lightning.callbacks.FinetuningScheduler
+  - class_path: pytorch_lightning.callbacks.finetuning_scheduler.FinetuningScheduler
   - class_path: pytorch_lightning.callbacks.finetuning_scheduler.FTSCheckpoint
     init_args:
         save_top_k: 5
         monitor: val_loss
         verbose: true
-  - class_path: pytorch_lightning.callbacks.EarlyStopping
+  - class_path: pytorch_lightning.callbacks.finetuning_scheduler.FTSEarlyStopping
     init_args:
       monitor: val_loss
       min_delta: 0.001

pl_examples/basic_examples/config/fts/nofts_baseline.yaml

@@ -1,6 +1,6 @@
 trainer:
   callbacks:
-  - class_path: pytorch_lightning.callbacks.EarlyStopping
+  - class_path: pytorch_lightning.callbacks.early_stopping.EarlyStopping
     init_args:
       monitor: val_loss
       min_delta: 0.001

pl_examples/basic_examples/fts_superglue.py

@@ -43,6 +43,7 @@
 
 import pytorch_lightning as pl
 from pl_examples import _HF_AVAILABLE
+from pytorch_lightning.callbacks.finetuning_scheduler.fts import FinetuningScheduler
 from pytorch_lightning.utilities import rank_zero_warn
 from pytorch_lightning.utilities.cli import instantiate_class, LightningCLI
 
@@ -116,6 +117,7 @@ def __init__(
         self.model.config.update(self.model_cfg)  # apply model config overrides
         self.metric = datasets.load_metric("super_glue", self.task_name, experiment_id=self.experiment_id)
         self.no_decay = ["bias", "LayerNorm.weight"]
+        self.finetuningscheduler_callback = None
 
     def forward(self, **inputs):
         return self.model(**inputs)
@@ -128,8 +130,8 @@ def training_step(self, batch, batch_idx):
     def training_epoch_end(self, outputs: List[Any]) -> None:
         loss = torch.stack([x["loss"] for x in outputs]).mean()
         self.log("train_loss", loss, prog_bar=True, sync_dist=True)
-        if self.trainer.finetuning_scheduler_callback:
-            self.log("finetuning_schedule_depth", self.trainer.finetuning_scheduler_callback.curr_depth)
+        if self.finetuningscheduler_callback:
+            self.log("finetuning_schedule_depth", self.finetuningscheduler_callback.curr_depth)
 
     def validation_step(self, batch, batch_idx, dataloader_idx=0):
         outputs = self(**batch)
@@ -183,6 +185,12 @@ def configure_optimizers(self):
         scheduler = {"scheduler": instantiate_class(optimizer, self.lr_scheduler_init), **self.pl_lrs_cfg}
         return [optimizer], [scheduler]
 
+    def configure_callbacks(self):
+        found_fts = [c for c in self.trainer.callbacks if isinstance(c, FinetuningScheduler)]
+        if found_fts:
+            self.finetuningscheduler_callback = found_fts[0]
+        return super().configure_callbacks()
+
 
 class RteBoolqDataModule(pl.LightningDataModule):
     """A :class:`~pytorch_lighting.core.LightningDataModule` for using either the RTE or BoolQ `SuperGLUE Hugging

pytorch_lightning/callbacks/__init__.py

@@ -15,7 +15,6 @@
 from pytorch_lightning.callbacks.device_stats_monitor import DeviceStatsMonitor
 from pytorch_lightning.callbacks.early_stopping import EarlyStopping
 from pytorch_lightning.callbacks.finetuning import BackboneFinetuning, BaseFinetuning
-from pytorch_lightning.callbacks.finetuning_scheduler import FinetuningScheduler  # idso
 from pytorch_lightning.callbacks.gpu_stats_monitor import GPUStatsMonitor
 from pytorch_lightning.callbacks.gradient_accumulation_scheduler import GradientAccumulationScheduler
 from pytorch_lightning.callbacks.lambda_function import LambdaCallback
@@ -32,11 +31,10 @@
 from pytorch_lightning.callbacks.xla_stats_monitor import XLAStatsMonitor
 
 __all__ = [
-    "FinetuningScheduler",
     "BackboneFinetuning",
     "BaseFinetuning",
-    "DeviceStatsMonitor",
     "Callback",
+    "DeviceStatsMonitor",
     "EarlyStopping",
     "GPUStatsMonitor",
     "XLAStatsMonitor",

pytorch_lightning/callbacks/early_stopping.py

@@ -112,8 +112,6 @@ def __init__(
         self.divergence_threshold = divergence_threshold
         self.wait_count = 0
         self.stopped_epoch = 0
-        self.es_phase_complete = True
-        self.final_phase = True
         self._check_on_train_epoch_end = check_on_train_epoch_end
 
         if self.mode not in self.mode_dict:
@@ -240,15 +238,12 @@ def _evaluate_stopping_criteria(self, current: torch.Tensor) -> Tuple[bool, Opti
         else:
             self.wait_count += 1
             if self.wait_count >= self.patience:
-                if self.final_phase:
-                    should_stop = True
-                    reason = (
-                        f"Monitored metric {self.monitor} did not improve in the last {self.wait_count} records."
-                        f" Best score: {self.best_score:.3f}. Signaling Trainer to stop."
-                    )
-                else:
-                    self.es_phase_complete = True
-                    self.wait_count = 0
+                should_stop = True
+                reason = (
+                    f"Monitored metric {self.monitor} did not improve in the last {self.wait_count} records."
+                    f" Best score: {self.best_score:.3f}. Signaling Trainer to stop."
+                )
+
         return should_stop, reason
 
     def _improvement_message(self, current: torch.Tensor) -> str:

pytorch_lightning/callbacks/finetuning_scheduler/__init__.py

@@ -21,7 +21,9 @@
 from pytorch_lightning.callbacks.finetuning_scheduler.fts_supporters import (  # noqa: F401 # isort: skip
     FTSState,
     FTSCheckpoint,
+    FTSEarlyStopping,
     SchedulingMixin,
     CallbackDepMixin,
+    CallbackResolverMixin,
 )
 from pytorch_lightning.callbacks.finetuning_scheduler.fts import FinetuningScheduler  # noqa: F401

pytorch_lightning/callbacks/finetuning_scheduler/fts.py

@@ -43,8 +43,10 @@ class FinetuningScheduler(BaseFinetuning, SchedulingMixin, CallbackDepMixin):
     unfreezing of models via a finetuning schedule that is either implicitly generated (the default) or explicitly
     provided by the user (more computationally efficient).
 
-    Finetuning phase transitions are driven by :class:`~pytorch_lightning.callbacks.early_stopping.EarlyStopping`
-    criteria, user-specified epoch transitions or a composition of the two (the default mode). A
+    Finetuning phase transitions are driven by
+    :class:`~pytorch_lightning.callbacks.finetuning_scheduler.fts_supporters.FTSEarlyStopping` criteria (a multi-phase
+    extension of :class:`~pytorch_lightning.callbacks.early_stopping.EarlyStopping`), user-specified epoch transitions
+    or a composition of the two (the default mode). A
     :class:`~pytorch_lightning.callbacks.finetuning_scheduler.fts.FinetuningScheduler` training session completes
     when the final phase of the schedule has its stopping criteria met. See
     :ref:`Early Stopping<common/early_stopping:Early stopping>` for more details on that callback's configuration.
@@ -103,11 +105,12 @@ def __init__(
                 and exit without training. Typically used to generate a default schedule that will be adjusted by the
                 user before training. Defaults to ``False``.
             epoch_transitions_only: If ``True``, Use epoch-driven stopping criteria exclusively (rather than composing
-                :class:`~pytorch_lightning.callbacks.early_stopping.EarlyStopping` and epoch-driven criteria which is
-                the default). If using this mode, an epoch-driven transition (``max_transition_epoch`` >= 0) must be
-                specified for each phase. If unspecified, ``max_transition_epoch`` defaults to -1 for each phase which
-                signals the application of :class:`~pytorch_lightning.callbacks.early_stopping.EarlyStopping` criteria
-                only. epoch_transitions_only defaults to ``False``.
+                :class:`~pytorch_lightning.callbacks.finetuning_scheduler.fts_supporters.FTSEarlyStopping` and
+                epoch-driven criteria which is the default). If using this mode, an epoch-driven transition
+                (``max_transition_epoch`` >= 0) must be specified for each phase. If unspecified,
+                ``max_transition_epoch`` defaults to -1 for each phase which signals the application of
+                :class:`~pytorch_lightning.callbacks.finetuning_scheduler.fts_supporters.FTSEarlyStopping` criteria only
+                . epoch_transitions_only defaults to ``False``.
 
         Attributes:
             _fts_state: The internal finetuning scheduler state.
@@ -268,7 +271,8 @@ def on_before_accelerator_backend_setup(self, trainer: "pl.Trainer", pl_module:
         """Before setting up the accelerator environment:
         Dump the default finetuning schedule
         OR
-        1. configure the :class:`~pytorch_lightning.callbacks.early_stopping.EarlyStopping` callback (if relevant)
+        1. configure the :class:`~pytorch_lightning.callbacks.finetuning_scheduler.fts_supporters.FTSEarlyStopping`
+            callback (if relevant)
         2. initialize the :attr:`~pytorch_lightning.callbacks.finetuning_scheduler.fts.FinetuningScheduler._fts_state`
         3. freeze the target :class:`~pytorch_lightning.core.lightning.LightningModule` parameters
 
@@ -386,8 +390,8 @@ def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -
 
     def should_transition(self, trainer: "pl.Trainer") -> bool:
         """Phase transition logic is contingent on whether we are composing
-        :class:`~pytorch_lightning.callbacks.early_stopping.EarlyStopping` criteria with epoch-driven transition
-        constraints or exclusively using epoch-driven transition scheduling. (i.e.,
+        :class:`~pytorch_lightning.callbacks.finetuning_scheduler.fts_supporters.FTSEarlyStopping` criteria with
+        epoch-driven transition constraints or exclusively using epoch-driven transition scheduling. (i.e.,
         :attr:`~pytorch_lightning.callbacks.finetuning_scheduler.fts.FinetuningScheduler.epoch_transitions_only` is
         ``True``)
 
@@ -400,7 +404,7 @@ def should_transition(self, trainer: "pl.Trainer") -> bool:
             if self.depth_remaining > 0
             else self.pl_module.trainer.fit_loop.max_epochs
         )
-        if not self.epoch_transitions_only:  # if we're considering EarlyStopping criteria
+        if not self.epoch_transitions_only:  # if we're considering FTSEarlyStopping criteria
             epoch_driven_transition = (
                 True
                 if not self.pl_module.trainer.early_stopping_callback.final_phase