drop LoggerStages

Author: Borda

pytorch_lightning/trainer/connectors/logger_connector/epoch_result_store.py

@@ -17,30 +17,10 @@
 import torch
 
 from pytorch_lightning.core.step_result import Result
+from pytorch_lightning.trainer.states import RunningStage
 from pytorch_lightning.utilities import DistributedType, LightningEnum
 
 
-class LoggerStages(LightningEnum):
-    """ Train/validation/test phase in each training step.
-
-    >>> # you can math the type with string
-    >>> LoggerStages.TRAIN == 'train'
-    True
-    """
-    TRAIN = "train"
-    VAL = "validation"
-    TEST = "test"
-
-    @staticmethod
-    def determine_stage(stage_or_testing: Union[str, bool]) -> 'LoggerStages':
-        if isinstance(stage_or_testing, str) and stage_or_testing in list(LoggerStages):
-            return LoggerStages(stage_or_testing)
-        if isinstance(stage_or_testing, (bool, int)):
-            # stage_or_testing is trainer.testing
-            return LoggerStages.TEST if bool(stage_or_testing) else LoggerStages.VAL
-        raise RuntimeError(f"Invalid stage {stage_or_testing} of type {type(stage_or_testing)} given")
-
-
 class ResultStoreType(LightningEnum):
     INSIDE_BATCH_TRAIN_LOOP = "inside_batch_train_loop"
     OUTSIDE_BATCH_TRAIN_LOOP = "outside_batch_train_loop"
@@ -276,7 +256,7 @@ class EpochResultStore:
 
     def __init__(self, trainer, stage):
         self.trainer = trainer
-        self._stage = stage
+        self._stage = RunningStage(stage)
         self.reset()
 
     def __getitem__(self, key: str) -> Any:
@@ -371,15 +351,14 @@ def update_logger_connector(self) -> None:
         callback_metrics = {}
         batch_pbar_metrics = {}
         batch_log_metrics = {}
-        is_train = self._stage in LoggerStages.TRAIN.value
 
         if not self._has_batch_loop_finished:
             # get pbar
             batch_pbar_metrics = self.get_latest_batch_pbar_metrics()
             logger_connector.add_progress_bar_metrics(batch_pbar_metrics)
             batch_log_metrics = self.get_latest_batch_log_metrics()
 
-            if is_train:
+            if self._stage == RunningStage.TRAINING:
                 # Only log and add to callback epoch step during evaluation, test.
                 logger_connector._logged_metrics.update(batch_log_metrics)
                 callback_metrics.update(batch_pbar_metrics)
@@ -401,7 +380,7 @@ def update_logger_connector(self) -> None:
             callback_metrics.update(epoch_log_metrics)
             callback_metrics.update(forked_metrics)
 
-        if not is_train and self.trainer.testing:
+        if self._stage != RunningStage.TRAINING and self.trainer.testing:
             logger_connector.evaluation_callback_metrics.update(callback_metrics)
 
         # update callback_metrics

pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py

@@ -22,8 +22,9 @@
 from pytorch_lightning.core.step_result import Result
 from pytorch_lightning.loggers import LoggerCollection, TensorBoardLogger
 from pytorch_lightning.trainer.connectors.logger_connector.callback_hook_validator import CallbackHookNameValidator
-from pytorch_lightning.trainer.connectors.logger_connector.epoch_result_store import EpochResultStore, LoggerStages
+from pytorch_lightning.trainer.connectors.logger_connector.epoch_result_store import EpochResultStore
 from pytorch_lightning.trainer.connectors.logger_connector.metrics_holder import MetricsHolder
+from pytorch_lightning.trainer.states import RunningStage
 from pytorch_lightning.utilities import DeviceType, flatten_dict
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.model_helpers import is_overridden
@@ -37,7 +38,7 @@ def __init__(self, trainer):
         self._logged_metrics = MetricsHolder()
         self._progress_bar_metrics = MetricsHolder()
         self.eval_loop_results = []
-        self._cached_results = {stage: EpochResultStore(trainer, stage) for stage in LoggerStages}
+        self._cached_results = {stage: EpochResultStore(trainer, stage) for stage in RunningStage}
         self._callback_hook_validator = CallbackHookNameValidator()
         self._current_stage = None
 
@@ -91,7 +92,10 @@ def set_metrics(self, key: str, val: Any) -> None:
         metrics_holder.reset(val)
 
     def set_stage(self, stage_or_testing: Union[str, bool], reset: bool = False) -> None:
-        self._current_stage = LoggerStages.determine_stage(stage_or_testing)
+        if isinstance(stage_or_testing, (bool, int)):
+            self._current_stage = RunningStage.TESTING if bool(stage_or_testing) else RunningStage.EVALUATING
+        else:
+            self._current_stage = RunningStage.from_str(stage_or_testing)
         if reset:
             self.cached_results.reset()