[2/N] Define dataclasses for progress tracking (#7574)

Co-authored-by: ananthsub <[email protected]>

Author: carmocca

CHANGELOG.md

@@ -21,7 +21,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added support for checkpointing based on a provided time interval during training ([#7515](https://github.com/PyTorchLightning/pytorch-lightning/pull/7515))
 
 
-- Added dataclasses for progress tracking ([#6603](https://github.com/PyTorchLightning/pytorch-lightning/pull/6603))
+- Added dataclasses for progress tracking (
+    [#6603](https://github.com/PyTorchLightning/pytorch-lightning/pull/6603),
+    [#7574](https://github.com/PyTorchLightning/pytorch-lightning/pull/7574))
 
 
 - Added argument `trainer.predict(ckpt_path)` ([#7430](https://github.com/PyTorchLightning/pytorch-lightning/pull/7430))

pytorch_lightning/trainer/progress.py

@@ -11,68 +11,97 @@
 # 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, field
+from typing import Optional
 
 
 @dataclass
-class ProgressState:
+class Tracker:
     """
-    Basic dataclass to track event progress.
+    Track an event's progress.
 
     Args:
         ready: Intended to track the number of events ready to start.
-        started: Intended to be incremented after the event is started (e.g. after `on_*_start runs).
+        started: Intended to be incremented after the event is started (e.g. after ``on_*_start`` runs).
         processed: Intended to be incremented after the event is processed.
-        completed: Intended to be incremented after the event completes (e.g. after `on_*_end` runs).
+        completed: Intended to be incremented after the event completes (e.g. after ``on_*_end`` runs).
+
+    Attributes set to ``None`` are treated as unused and are restricted.
     """
-    ready: int = 0
-    started: int = 0
-    processed: int = 0
-    completed: int = 0
+    ready: Optional[int] = 0
+    started: Optional[int] = 0
+    processed: Optional[int] = 0
+    completed: Optional[int] = 0
 
     def reset(self) -> None:
-        self.ready = 0
-        self.started = 0
-        self.processed = 0
-        self.completed = 0
+        if self.ready is not None:
+            self.ready = 0
+        if self.started is not None:
+            self.started = 0
+        if self.processed is not None:
+            self.processed = 0
+        if self.completed is not None:
+            self.completed = 0
+
+    def __setattr__(self, key: str, value: int) -> None:
+        if getattr(self, key, 0) is None:
+            raise AttributeError(f"The '{key}' attribute is meant to be unused")
+        return super().__setattr__(key, value)
+
+    def __repr__(self):
+        # hide `None` fields
+        args = [f"{k}={v}" for k, v in self.__dict__.items() if v is not None]
+        return f"{self.__class__.__name__}({', '.join(args)})"
 
 
 @dataclass
 class Progress:
     """
-    Basic dataclass to track aggregated and current progress states.
+    Track aggregated and current progress.
 
     Args:
         total: Intended to track the total progress of an event
         current: Intended to track the current progress of an event
     """
-    total: ProgressState = field(default_factory=ProgressState)
-    current: ProgressState = field(default_factory=ProgressState)
+    total: Tracker = field(default_factory=Tracker)
+    current: Tracker = field(default_factory=Tracker)
 
     def increment_ready(self) -> None:
+        if self.total.ready is None or self.current.ready is None:
+            return
         self.total.ready += 1
         self.current.ready += 1
 
     def increment_started(self) -> None:
+        if self.total.started is None or self.current.started is None:
+            return
         self.total.started += 1
         self.current.started += 1
 
     def increment_processed(self) -> None:
+        if self.total.processed is None or self.current.processed is None:
+            return
         self.total.processed += 1
         self.current.processed += 1
 
     def increment_completed(self) -> None:
+        if self.total.completed is None or self.current.completed is None:
+            return
         self.total.completed += 1
         self.current.completed += 1
 
+    @classmethod
+    def from_defaults(cls, **kwargs: Optional[int]) -> 'Progress':
+        return cls(total=Tracker(**kwargs), current=Tracker(**kwargs))
+
 
 @dataclass
 class LoopProgress:
     """
-    Dataclass to track loop progress during execution.
+    Track loop progress during execution.
 
     These counters are local to a trainer rank. By default, they are not globally synced across all ranks.
+
     Args:
         epoch: Tracks epochs progress.
         batch: Tracks batch progress.
@@ -87,3 +116,65 @@ def increment_epoch_completed(self) -> None:
     def reset_on_epoch(self) -> None:
         self.batch.current.reset()
         self.epoch.current.reset()
+
+
+@dataclass
+class OptimizationProgress:
+    """
+    Track optimization progress.
+
+    Args:
+        optimizer: Tracks optimizer progress.
+        scheduler: Tracks scheduler progress.
+    """
+    optimizer: Progress = Progress.from_defaults(processed=None)
+    scheduler: Progress = Progress.from_defaults(started=None, processed=None)
+    zero_grad: Progress = Progress.from_defaults(processed=None)
+
+    @property
+    def optimizer_steps(self) -> int:
+        return self.optimizer.total.completed
+
+    @property
+    def scheduler_steps(self) -> int:
+        return self.scheduler.total.completed
+
+
+@dataclass
+class TrainingProgress(Progress):
+    """
+    Extends ``Progress`` with training specific attributes
+
+    Args:
+        optimization: Tracks optimization progress
+    """
+    optimization: OptimizationProgress = field(default_factory=OptimizationProgress)
+
+
+@dataclass
+class TrainingLoopProgress(LoopProgress):
+    epoch: TrainingProgress = field(default_factory=TrainingProgress)
+
+    def reset_on_epoch(self) -> None:
+        # override to avoid resetting `epoch.current`
+        self.batch.current.reset()
+
+
+@dataclass
+class FitLoopProgress:
+    train: TrainingLoopProgress = field(default_factory=TrainingLoopProgress)
+    val: LoopProgress = field(default_factory=LoopProgress)
+
+
+@dataclass
+class LoopState:
+    """
+    Basic dataclass to track loop progress across trainer functions during trainer execution.
+
+    This class will be removed and these attributes will live in each loop.
+    """
+
+    fit: FitLoopProgress = field(default_factory=FitLoopProgress)
+    val: LoopProgress = field(default_factory=LoopProgress)
+    test: LoopProgress = field(default_factory=LoopProgress)
+    predict: LoopProgress = field(default_factory=LoopProgress)