File size: 2,891 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
# Copyright The Lightning AI team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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 dataclasses import dataclass
from typing import Optional
from pytorch_lightning.utilities.enums import LightningEnum
class TrainerStatus(LightningEnum):
"""Enum for the status of the :class:`~pytorch_lightning.trainer.trainer.Trainer`"""
INITIALIZING = "initializing" # trainer creation
RUNNING = "running"
FINISHED = "finished"
INTERRUPTED = "interrupted"
@property
def stopped(self) -> bool:
return self in (self.FINISHED, self.INTERRUPTED)
class TrainerFn(LightningEnum):
"""Enum for the user-facing functions of the :class:`~pytorch_lightning.trainer.trainer.Trainer` such as
:meth:`~pytorch_lightning.trainer.trainer.Trainer.fit` and
:meth:`~pytorch_lightning.trainer.trainer.Trainer.test`."""
FITTING = "fit"
VALIDATING = "validate"
TESTING = "test"
PREDICTING = "predict"
class RunningStage(LightningEnum):
"""Enum for the current running stage.
This stage complements :class:`TrainerFn` by specifying the current running stage for each function.
More than one running stage value can be set while a :class:`TrainerFn` is running:
- ``TrainerFn.FITTING`` - ``RunningStage.{SANITY_CHECKING,TRAINING,VALIDATING}``
- ``TrainerFn.VALIDATING`` - ``RunningStage.VALIDATING``
- ``TrainerFn.TESTING`` - ``RunningStage.TESTING``
- ``TrainerFn.PREDICTING`` - ``RunningStage.PREDICTING``
"""
TRAINING = "train"
SANITY_CHECKING = "sanity_check"
VALIDATING = "validate"
TESTING = "test"
PREDICTING = "predict"
@property
def evaluating(self) -> bool:
return self in (self.VALIDATING, self.TESTING, self.SANITY_CHECKING)
@property
def dataloader_prefix(self) -> Optional[str]:
if self in (self.VALIDATING, self.SANITY_CHECKING):
return "val"
return self.value
@dataclass
class TrainerState:
"""Dataclass to encapsulate the current :class:`~pytorch_lightning.trainer.trainer.Trainer` state."""
status: TrainerStatus = TrainerStatus.INITIALIZING
fn: Optional[TrainerFn] = None
stage: Optional[RunningStage] = None
@property
def finished(self) -> bool:
return self.status == TrainerStatus.FINISHED
@property
def stopped(self) -> bool:
return self.status.stopped
|