[1/N] Define dataclasses for progress tracking (#6603)

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

Author: ananthsub

CHANGELOG.md

@@ -14,6 +14,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added LightningCLI support for config files on object stores ([#7521](https://github.com/PyTorchLightning/pytorch-lightning/pull/7521))
 
 
+- Added dataclasses for progress tracking ([#6603](https://github.com/PyTorchLightning/pytorch-lightning/pull/6603))
+
+
 - Added argument `trainer.predict(ckpt_path)` ([#7430](https://github.com/PyTorchLightning/pytorch-lightning/pull/7430))
 
 

pytorch_lightning/trainer/progress.py

@@ -0,0 +1,89 @@
+# Copyright The PyTorch Lightning 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, field
+
+
+@dataclass
+class ProgressState:
+    """
+    Basic dataclass to track event 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).
+        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).
+    """
+    ready: int = 0
+    started: int = 0
+    processed: int = 0
+    completed: int = 0
+
+    def reset(self) -> None:
+        self.ready = 0
+        self.started = 0
+        self.processed = 0
+        self.completed = 0
+
+
+@dataclass
+class Progress:
+    """
+    Basic dataclass to track aggregated and current progress states.
+
+    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)
+
+    def increment_ready(self) -> None:
+        self.total.ready += 1
+        self.current.ready += 1
+
+    def increment_started(self) -> None:
+        self.total.started += 1
+        self.current.started += 1
+
+    def increment_processed(self) -> None:
+        self.total.processed += 1
+        self.current.processed += 1
+
+    def increment_completed(self) -> None:
+        self.total.completed += 1
+        self.current.completed += 1
+
+
+@dataclass
+class LoopProgress:
+    """
+    Dataclass to 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.
+    """
+    epoch: Progress = field(default_factory=Progress)
+    batch: Progress = field(default_factory=Progress)
+
+    def increment_epoch_completed(self) -> None:
+        self.epoch.increment_completed()
+        self.reset_on_epoch()
+
+    def reset_on_epoch(self) -> None:
+        self.batch.current.reset()
+        self.epoch.current.reset()

tests/trainer/test_progress.py

@@ -0,0 +1,150 @@
+# Copyright The PyTorch Lightning 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 pytorch_lightning.trainer.progress import LoopProgress, ProgressState
+
+
+def test_increment_ready(tmpdir):
+    prog = LoopProgress()
+    prog.batch.increment_ready()
+    assert prog.batch.total.ready == 1
+    assert prog.batch.current.ready == prog.batch.total.ready
+
+
+def test_increment_started(tmpdir):
+    prog = LoopProgress()
+    prog.epoch.increment_started()
+    assert prog.batch.total.ready == 0
+    assert prog.epoch.total.ready == 0
+    assert prog.epoch.total.started == 1
+    assert prog.epoch.total.started == prog.epoch.current.started
+
+
+def test_increment_processed(tmpdir):
+    prog = LoopProgress()
+    prog.epoch.increment_processed()
+    assert prog.batch.total.ready == 0
+    assert prog.batch.total.started == 0
+    assert prog.epoch.total.started == 0
+    assert prog.epoch.total.processed == 1
+    assert prog.epoch.total.processed == prog.epoch.current.processed
+
+
+def test_increment_completed(tmpdir):
+    prog = LoopProgress()
+    prog.epoch.increment_completed()
+    assert prog.batch.total.ready == 0
+    assert prog.batch.total.started == 0
+    assert prog.epoch.total.started == 0
+    assert prog.epoch.total.processed == 0
+    assert prog.epoch.total.completed == 1
+    assert prog.epoch.total.completed == prog.epoch.current.completed
+
+
+def test_increment_epoch(tmpdir):
+    """ Test sequences for incrementing epochs. """
+    prog = LoopProgress()
+    prog.batch.increment_completed()
+    assert prog.batch.current.completed == 1
+
+    prog.increment_epoch_completed()
+    prog.increment_epoch_completed()
+    assert prog.epoch.current.completed == 0
+    assert prog.epoch.total.completed == 2
+    assert prog.batch.current.completed == 0
+    assert prog.batch.total.completed == 1
+
+
+def test_reset_on_epoch(tmpdir):
+    """ Test sequences for resetting. """
+    prog = LoopProgress()
+
+    prog.batch.increment_started()
+    assert prog.batch.total.started == 1
+    assert prog.epoch.total.started == 0
+
+    prog.reset_on_epoch()
+    assert prog.batch.current.started == 0
+    assert prog.batch.total == ProgressState(started=1)
+
+    prog.batch.increment_started()
+    assert prog.batch.total == ProgressState(started=2)
+    assert prog.epoch.total.started == 0
+
+
+def test_increment_batch_ready_start_process_finish_epoch(tmpdir):
+    """ Test sequences for incrementing batches reads and epochs. """
+    prog = LoopProgress()
+
+    prog.epoch.increment_ready()
+    assert prog.epoch.total.ready == 1
+    assert prog.epoch.current.ready == 1
+    assert prog.batch.total.ready == 0
+    assert prog.batch.current.ready == 0
+
+    prog.epoch.increment_started()
+    assert prog.epoch.total.ready == 1
+    assert prog.epoch.current.ready == 1
+    assert prog.epoch.total.started == 1
+    assert prog.epoch.current.started == 1
+    assert prog.batch.total.started == 0
+    assert prog.batch.current.started == 0
+
+    prog.batch.increment_ready()
+    assert prog.batch.total.ready == 1
+    assert prog.batch.current.ready == 1
+    assert prog.epoch.total.ready == 1
+    assert prog.epoch.current.ready == 1
+
+    prog.batch.increment_started()
+    assert prog.batch.total.started == 1
+    assert prog.batch.current.started == 1
+    assert prog.epoch.total.started == 1
+    assert prog.epoch.current.started == 1
+
+    prog.batch.increment_processed()
+    assert prog.batch.total.processed == 1
+    assert prog.batch.current.processed == 1
+    assert prog.epoch.total.processed == 0
+    assert prog.epoch.current.processed == 0
+
+    prog.batch.increment_completed()
+    assert prog.batch.total.completed == 1
+    assert prog.batch.current.completed == 1
+    assert prog.epoch.total.completed == 0
+    assert prog.epoch.current.completed == 0
+
+    prog.epoch.increment_processed()
+    assert prog.batch.total.processed == 1
+    assert prog.batch.current.processed == 1
+    assert prog.epoch.total.processed == 1
+    assert prog.epoch.current.processed == 1
+
+    prog.increment_epoch_completed()
+    assert prog.batch.total.completed == 1
+    assert prog.batch.current.completed == 0
+    assert prog.epoch.total.completed == 1
+    assert prog.epoch.current.completed == 0
+
+    prog.epoch.increment_ready()
+    assert prog.epoch.total.ready == 2
+    assert prog.epoch.current.ready == 1
+
+    prog.batch.increment_ready()
+    assert prog.batch.total.ready == 2
+    assert prog.batch.current.ready == 1
+
+    prog.reset_on_epoch()
+    assert prog.batch.current.ready == 0
+    assert prog.epoch.current.ready == 0