simplify training phase as Enum (#5419)

* simplify training phase as Enum

* tests

* .

* .

* rename

* Apply suggestions from code review

Co-authored-by: Carlos Mocholí <[email protected]>

* Apply suggestions from code review

Co-authored-by: Carlos Mocholí <[email protected]>

* rename

* flake8

Co-authored-by: Carlos Mocholí <[email protected]>

Author: Borda

pytorch_lightning/trainer/deprecated_api.py

@@ -11,14 +11,15 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+from pytorch_lightning.trainer.states import RunningStage
 from pytorch_lightning.utilities import DistributedType, DeviceType, rank_zero_warn
 
 
 class DeprecatedDistDeviceAttributes:
 
     _distrib_type: DistributedType
     _device_type: DeviceType
+    _running_stage: RunningStage
     num_gpus: int
 
     @property
@@ -129,3 +130,27 @@ def use_single_gpu(self, val: bool) -> None:
         )
         if val:
             self._device_type = DeviceType.GPU
+
+    @property
+    def training(self) -> bool:
+        # todo: consider rename as `is_training`
+        return self._running_stage == RunningStage.TRAINING
+
+    @training.setter
+    def training(self, val: bool) -> None:
+        if val:
+            self._running_stage = RunningStage.TRAINING
+        else:
+            self._running_stage = None
+
+    @property
+    def testing(self) -> bool:
+        # todo: consider rename as `is_testing`
+        return self._running_stage == RunningStage.TESTING
+
+    @testing.setter
+    def testing(self, val: bool) -> None:
+        if val:
+            self._running_stage = RunningStage.TESTING
+        else:
+            self._running_stage = None

pytorch_lightning/trainer/states.py

@@ -23,10 +23,10 @@ class TrainerState(LightningEnum):
     """ State which is set in the :class:`~pytorch_lightning.trainer.trainer.Trainer`
     to indicate what is currently or was executed.
 
-    >>> # you can math the type with string
+    >>> # you can compare the type with a string
     >>> TrainerState.RUNNING == 'RUNNING'
     True
-    >>> # which is case sensitive
+    >>> # which is case insensitive
     >>> TrainerState.FINISHED == 'finished'
     True
     """
@@ -36,6 +36,19 @@ class TrainerState(LightningEnum):
     INTERRUPTED = 'INTERRUPTED'
 
 
+class RunningStage(LightningEnum):
+    """Type of train phase.
+
+    >>> # you can match the Enum with string
+    >>> RunningStage.TRAINING == 'train'
+    True
+    """
+    TRAINING = 'train'
+    EVALUATING = 'eval'
+    TESTING = 'test'
+    TUNING = 'tune'
+
+
 def trainer_state(*, entering: Optional[TrainerState] = None, exiting: Optional[TrainerState] = None) -> Callable:
     """ Decorator for :class:`~pytorch_lightning.trainer.trainer.Trainer` methods
     which changes state to `entering` before the function execution and `exiting`

pytorch_lightning/trainer/trainer.py

@@ -294,6 +294,7 @@ def __init__(
         super().__init__()
         self._device_type = DeviceType.CPU
         self._distrib_type = None
+        self._running_stage = None
 
         # init connectors
         self.dev_debugger = InternalDebugger(self)

tests/deprecated_api/test_remove_1-4.py

@@ -42,7 +42,7 @@ def test_v1_4_0_deprecated_imports():
         from pytorch_lightning.utilities.xla_device_utils import XLADeviceUtils  # noqa: F811 F401
 
 
-def test_v1_4_0_deprecated_trainer_attributes():
+def test_v1_4_0_deprecated_trainer_device_distrib():
     """Test that Trainer attributes works fine."""
     trainer = Trainer()
     trainer._distrib_type = None
@@ -89,6 +89,22 @@ def test_v1_4_0_deprecated_trainer_attributes():
         assert trainer.use_horovod
 
 
+def test_v1_4_0_deprecated_trainer_phase():
+    """Test that Trainer attributes works fine."""
+    trainer = Trainer()
+
+    assert not trainer.training
+    assert not trainer.testing
+
+    trainer.training = True
+    assert trainer.training
+    assert not trainer.testing
+
+    trainer.testing = True
+    assert not trainer.training
+    assert trainer.testing
+
+
 def test_v1_4_0_deprecated_metrics():
     from pytorch_lightning.metrics.functional.classification import stat_scores_multiple_classes
     with pytest.deprecated_call(match='will be removed in v1.4'):