Breaking Change - Move progress to Unit (#459)

Summary:
Pull Request resolved: #459

1. Progress is no longer tracked on the State but on the Unit
2. Removed the `get_current_progress` utility

Differential Revision: D47358509

fbshipit-source-id: 3fa1836286d6cd89e8c73cea28dd789322ea75cc

examples/auto_unit_example.py

@@ -18,7 +18,7 @@
 from torch.utils.data.dataset import Dataset, TensorDataset
 from torcheval.metrics import BinaryAccuracy
 from torchtnt.framework import AutoUnit, fit, init_fit_state, State
-from torchtnt.framework.utils import get_current_progress
+from torchtnt.framework.state import EntryPoint
 from torchtnt.utils import init_from_env, seed, TLRScheduler
 from torchtnt.utils.loggers import TensorBoardLogger
 
@@ -120,7 +120,10 @@ def on_eval_step_end(
         self.eval_accuracy.update(outputs, targets)
 
     def on_eval_end(self, state: State) -> None:
-        step = get_current_progress(state).num_steps_completed
+        if state.entry_point == EntryPoint.FIT:
+            step = self.train_progress.num_steps_completed
+        else:
+            step = self.eval_progress.num_steps_completed
         accuracy = self.eval_accuracy.compute()
         self.tb_logger.log("eval_accuracy", accuracy, step)
         self.eval_accuracy.reset()

examples/mnist/main.py

@@ -106,12 +106,6 @@ def on_train_epoch_end(self, state: State) -> None:
     def on_eval_step_end(
         self, state: State, data: Batch, step: int, loss: torch.Tensor, outputs: Any
     ) -> None:
-        # step_count = state.eval_state.progress.num_steps_completed
-        # data = copy_data_to_device(data, self.device)
-        # inputs, targets = data
-
-        # outputs = self.module(inputs)
-        # loss = torch.nn.functional.nll_loss(outputs, targets)
         if step % self.log_every_n_steps == 0:
             self.tb_logger.log("evaluation loss", loss, step)
 

examples/torchdata_train_example.py

@@ -115,15 +115,15 @@ def train_step(self, state: State, data: Batch) -> None:
 
         # update metrics & logs
         self.train_accuracy.update(outputs, targets)
-        step_count = state.train_state.progress.num_steps_completed
+        step_count = self.train_progress.num_steps_completed
         if (step_count + 1) % self.log_every_n_steps == 0:
             accuracy = self.train_accuracy.compute()
             self.tb_logger.log("loss", loss, step_count)
             self.tb_logger.log("accuracy", accuracy, step_count)
 
     def on_train_epoch_end(self, state: State) -> None:
         # compute and log the metrics at the end of epoch
-        step_count = state.train_state.progress.num_steps_completed
+        step_count = self.train_progress.num_steps_completed
         accuracy = self.train_accuracy.compute()
         self.tb_logger.log("accuracy_epoch", accuracy, step_count)
 

examples/torchrec/main.py

@@ -42,7 +42,6 @@
 
 from torchtnt.framework import EvalUnit, fit, init_fit_state, State, TrainUnit
 from torchtnt.framework.callbacks import TQDMProgressBar
-from torchtnt.framework.utils import get_current_progress
 from torchtnt.utils import (
     get_process_group_backend_from_device,
     init_from_env,
@@ -202,7 +201,7 @@ def __init__(
         self.log_every_n_steps = log_every_n_steps
 
     def train_step(self, state: State, data: Iterator[Batch]) -> None:
-        step = get_current_progress(state).num_steps_completed
+        step = self.train_progress.num_steps_completed
         loss, logits, labels = self.pipeline.progress(data)
         preds = torch.sigmoid(logits)
         self.train_auroc.update(preds, labels)
@@ -217,7 +216,7 @@ def on_train_epoch_end(self, state: State) -> None:
         self.train_auroc.reset()
 
     def eval_step(self, state: State, data: Iterator[Batch]) -> None:
-        step = get_current_progress(state).num_steps_completed
+        step = self.eval_progress.num_steps_completed
         loss, _, _ = self.pipeline.progress(data)
         if step % self.log_every_n_steps == 0:
             self.tb_logger.log("evaluation_loss", loss, step)

examples/train_unit_example.py

@@ -93,15 +93,15 @@ def train_step(self, state: State, data: Batch) -> None:
 
         # update metrics & logs
         self.train_accuracy.update(outputs, targets)
-        step_count = state.train_state.progress.num_steps_completed
+        step_count = self.train_progress.num_steps_completed
         if (step_count + 1) % self.log_every_n_steps == 0:
             acc = self.train_accuracy.compute()
             self.tb_logger.log("loss", loss, step_count)
             self.tb_logger.log("accuracy", acc, step_count)
 
     def on_train_epoch_end(self, state: State) -> None:
         # compute and log the metric at the end of the epoch
-        step_count = state.train_state.progress.num_steps_completed
+        step_count = self.train_progress.num_steps_completed
         acc = self.train_accuracy.compute()
         self.tb_logger.log("accuracy_epoch", acc, step_count)
 

tests/framework/callbacks/test_csv_writer.py

@@ -49,7 +49,7 @@ def test_csv_writer(self) -> None:
         dataset_len = 10
         batch_size = 2
 
-        my_unit = MagicMock(spec=DummyPredictUnit)
+        my_unit = DummyPredictUnit(2)
         dataloader = generate_random_dataloader(dataset_len, input_dim, batch_size)
         state = init_predict_state(dataloader=dataloader)
 
@@ -73,7 +73,7 @@ def test_csv_writer_single_row(self) -> None:
         dataset_len = 10
         batch_size = 2
 
-        my_unit = MagicMock(spec=DummyPredictUnit)
+        my_unit = DummyPredictUnit(2)
         dataloader = generate_random_dataloader(dataset_len, input_dim, batch_size)
         state = init_predict_state(dataloader=dataloader)
 
@@ -96,7 +96,7 @@ def test_csv_writer_with_no_output_rows_def(self) -> None:
         dataset_len = 10
         batch_size = 2
 
-        my_unit = MagicMock(spec=DummyPredictUnit)
+        my_unit = DummyPredictUnit(2)
         dataloader = generate_random_dataloader(dataset_len, input_dim, batch_size)
         state = init_predict_state(dataloader=dataloader)
 

tests/framework/callbacks/test_garbage_collector.py

@@ -35,7 +35,7 @@ def test_garbage_collector_call_count_train(self) -> None:
         max_epochs = 2
         expected_num_total_steps = dataset_len / batch_size * max_epochs
 
-        my_unit = MagicMock(spec=DummyTrainUnit)
+        my_unit = DummyTrainUnit(2)
         gc_callback_mock = MagicMock(spec=GarbageCollector)
 
         dataloader = generate_random_dataloader(dataset_len, input_dim, batch_size)
@@ -57,7 +57,7 @@ def test_garbage_collector_enabled_train(self) -> None:
         batch_size = 2
         max_epochs = 2
 
-        my_unit = MagicMock(spec=DummyTrainUnit)
+        my_unit = DummyTrainUnit(2)
         gc_callback = GarbageCollector(2)
 
         dataloader = generate_random_dataloader(dataset_len, input_dim, batch_size)
@@ -76,7 +76,7 @@ def test_garbage_collector_call_count_evaluate(self) -> None:
         batch_size = 2
         expected_num_total_steps = dataset_len / batch_size
 
-        my_unit = MagicMock(spec=DummyEvalUnit)
+        my_unit = DummyEvalUnit(2)
         gc_callback_mock = MagicMock(spec=GarbageCollector)
 
         dataloader = generate_random_dataloader(dataset_len, input_dim, batch_size)
@@ -97,7 +97,7 @@ def test_garbage_collector_enabled_evaluate(self) -> None:
         dataset_len = 10
         batch_size = 2
 
-        my_unit = MagicMock(spec=DummyEvalUnit)
+        my_unit = DummyEvalUnit(2)
         gc_callback = GarbageCollector(2)
 
         dataloader = generate_random_dataloader(dataset_len, input_dim, batch_size)
@@ -116,7 +116,7 @@ def test_garbage_collector_call_count_predict(self) -> None:
         batch_size = 2
         expected_num_total_steps = dataset_len / batch_size
 
-        my_unit = MagicMock(spec=DummyPredictUnit)
+        my_unit = DummyPredictUnit(2)
         gc_callback_mock = MagicMock(spec=GarbageCollector)
 
         dataloader = generate_random_dataloader(dataset_len, input_dim, batch_size)
@@ -137,7 +137,7 @@ def test_garbage_collector_enabled_predict(self) -> None:
         dataset_len = 10
         batch_size = 2
 
-        my_unit = MagicMock(spec=DummyPredictUnit)
+        my_unit = DummyPredictUnit(2)
         gc_callback = GarbageCollector(2)
 
         dataloader = generate_random_dataloader(dataset_len, input_dim, batch_size)
@@ -163,7 +163,7 @@ def test_garbage_collector_call_count_fit(self) -> None:
         )
         gc_step_interval = 4
 
-        my_unit = MagicMock(spec=DummyFitUnit)
+        my_unit = DummyFitUnit(2)
         gc_callback = GarbageCollector(gc_step_interval)
 
         train_dataloader = generate_random_dataloader(
@@ -201,7 +201,7 @@ def test_garbage_collector_enabled_fit(self) -> None:
         max_epochs = 2
         evaluate_every_n_epochs = 1
 
-        my_unit = MagicMock(spec=DummyFitUnit)
+        my_unit = DummyFitUnit(2)
         gc_callback = GarbageCollector(2)
 
         train_dataloader = generate_random_dataloader(

tests/framework/callbacks/test_pytorch_profiler.py

@@ -53,7 +53,7 @@ def test_profiler_evaluate(self) -> None:
         batch_size = 2
         expected_num_total_steps = dataset_len / batch_size
 
-        my_unit = MagicMock(spec=DummyEvalUnit)
+        my_unit = DummyEvalUnit(2)
         profiler_mock = MagicMock(spec=torch.profiler.profile)
 
         profiler = PyTorchProfiler(profiler=profiler_mock)
@@ -75,7 +75,7 @@ def test_profiler_predict(self) -> None:
         batch_size = 2
         expected_num_total_steps = dataset_len / batch_size
 
-        my_unit = MagicMock(spec=DummyPredictUnit)
+        my_unit = DummyPredictUnit(2)
         profiler_mock = MagicMock(spec=torch.profiler.profile)
 
         profiler = PyTorchProfiler(profiler=profiler_mock)

tests/framework/callbacks/test_torchsnapshot_saver.py

@@ -122,12 +122,10 @@ def test_save_restore(self) -> None:
             )
             train(state, my_unit, callbacks=[snapshot_cb])
 
-            end_num_steps_completed = state.train_state.progress.num_steps_completed
+            end_num_steps_completed = my_unit.train_progress.num_steps_completed
             self.assertGreater(len(expected_paths), 0)
             snapshot_cb.restore(expected_paths[0], state, my_unit)
-            restored_num_steps_completed = (
-                state.train_state.progress.num_steps_completed
-            )
+            restored_num_steps_completed = my_unit.train_progress.num_steps_completed
             # A snapshot is saved every n steps
             # so the first snapshot's progress will be equal to save_every_n_train_steps
             self.assertNotEqual(restored_num_steps_completed, end_num_steps_completed)

tests/framework/callbacks/test_tqdm_progress_bar.py

@@ -6,7 +6,6 @@
 # LICENSE file in the root directory of this source tree.
 
 import unittest
-from unittest.mock import MagicMock
 
 from torchtnt.framework._test_utils import (
     DummyEvalUnit,
@@ -39,7 +38,7 @@ def test_progress_bar_train(self) -> None:
             ),
         )
 
-        my_unit = MagicMock(spec=DummyTrainUnit)
+        my_unit = DummyTrainUnit(2)
         progress_bar = TQDMProgressBar()
         progress_bar.on_train_epoch_start(state, my_unit)
         self.assertEqual(progress_bar._train_progress_bar.total, expected_total)
@@ -56,7 +55,7 @@ def test_progress_bar_train_integration(self) -> None:
         dataloader = generate_random_dataloader(dataset_len, input_dim, batch_size)
         state = init_train_state(dataloader=dataloader, max_epochs=max_epochs)
 
-        my_unit = MagicMock(spec=DummyTrainUnit)
+        my_unit = DummyTrainUnit(2)
         progress_bar = TQDMProgressBar()
         train(state, my_unit, callbacks=[progress_bar])
 
@@ -79,7 +78,7 @@ def test_progress_bar_evaluate(self) -> None:
             ),
         )
 
-        my_unit = MagicMock(spec=DummyEvalUnit)
+        my_unit = DummyEvalUnit(2)
         progress_bar = TQDMProgressBar()
         progress_bar.on_eval_epoch_start(state, my_unit)
         self.assertEqual(progress_bar._eval_progress_bar.total, expected_total)
@@ -103,7 +102,7 @@ def test_progress_bar_predict(self) -> None:
             ),
         )
 
-        my_unit = MagicMock(spec=DummyPredictUnit)
+        my_unit = DummyPredictUnit(2)
         progress_bar = TQDMProgressBar()
         progress_bar.on_predict_epoch_start(state, my_unit)
         self.assertEqual(progress_bar._predict_progress_bar.total, expected_total)
@@ -126,9 +125,8 @@ def test_progress_bar_mid_progress(self) -> None:
                 max_epochs=max_epochs,
             ),
         )
-        state.predict_state.progress._num_steps_completed = 2
-
-        my_unit = MagicMock(spec=DummyPredictUnit)
+        my_unit = DummyPredictUnit(2)
+        my_unit.predict_progress._num_steps_completed = 2
         progress_bar = TQDMProgressBar()
         progress_bar.on_predict_epoch_start(state, my_unit)
         self.assertEqual(progress_bar._predict_progress_bar.total, expected_total)

tests/framework/test_auto_unit.py

@@ -465,7 +465,7 @@ def _test_ddp_no_sync() -> None:
             auto_unit.train_step(state=state, data=dummy_iterator)
             no_sync_mock.assert_called_once()
 
-        state.train_state.progress.increment_step()
+        auto_unit.train_progress.increment_step()
         # for the second step no_sync should not be called since we run optimizer step
         with patch.object(auto_unit.module, "no_sync") as no_sync_mock:
             auto_unit.train_step(state=state, data=dummy_iterator)
@@ -496,7 +496,7 @@ def _test_fsdp_no_sync() -> None:
             auto_unit.train_step(state=state, data=dummy_iterator)
             no_sync_mock.assert_called_once()
 
-        state.train_state.progress.increment_step()
+        auto_unit.train_progress.increment_step()
         # for the second step no_sync should not be called since we run optimizer step
         with patch.object(auto_unit.module, "no_sync") as no_sync_mock:
             auto_unit.train_step(state=state, data=dummy_iterator)
@@ -1108,7 +1108,7 @@ def compute_loss(
         tc = unittest.TestCase()
         tc.assertEqual(
             self._is_last_train_batch,
-            state.train_state.progress.num_steps_completed_in_epoch + 1
+            self.train_progress.num_steps_completed_in_epoch + 1
             == self.expected_steps_per_epoch,
         )
         inputs, targets = data

tests/framework/test_evaluate.py

@@ -34,9 +34,9 @@ def test_evaluate(self) -> None:
         state = init_eval_state(dataloader=dataloader)
         evaluate(state, my_unit)
 
-        self.assertEqual(state.eval_state.progress.num_epochs_completed, 1)
-        self.assertEqual(state.eval_state.progress.num_steps_completed_in_epoch, 0)
-        self.assertEqual(state.eval_state.progress.num_steps_completed, expected_steps)
+        self.assertEqual(my_unit.eval_progress.num_epochs_completed, 1)
+        self.assertEqual(my_unit.eval_progress.num_steps_completed_in_epoch, 0)
+        self.assertEqual(my_unit.eval_progress.num_steps_completed, expected_steps)
         self.assertEqual(state.entry_point, EntryPoint.EVALUATE)
 
         # step_output should be reset to None
@@ -62,11 +62,9 @@ def test_evaluate_max_steps_per_epoch(self) -> None:
         )
         evaluate(state, my_unit)
 
-        self.assertEqual(state.eval_state.progress.num_epochs_completed, 1)
-        self.assertEqual(state.eval_state.progress.num_steps_completed_in_epoch, 0)
-        self.assertEqual(
-            state.eval_state.progress.num_steps_completed, max_steps_per_epoch
-        )
+        self.assertEqual(my_unit.eval_progress.num_epochs_completed, 1)
+        self.assertEqual(my_unit.eval_progress.num_steps_completed_in_epoch, 0)
+        self.assertEqual(my_unit.eval_progress.num_steps_completed, max_steps_per_epoch)
         self.assertEqual(state.entry_point, EntryPoint.EVALUATE)
 
         # step_output should be reset to None
@@ -93,10 +91,10 @@ def test_evaluate_stop(self) -> None:
         )
         evaluate(state, my_unit)
 
-        self.assertEqual(state.eval_state.progress.num_epochs_completed, 1)
-        self.assertEqual(state.eval_state.progress.num_steps_completed_in_epoch, 0)
+        self.assertEqual(my_unit.eval_progress.num_epochs_completed, 1)
+        self.assertEqual(my_unit.eval_progress.num_steps_completed_in_epoch, 0)
         self.assertEqual(
-            my_unit.steps_processed, state.eval_state.progress.num_steps_completed
+            my_unit.steps_processed, my_unit.eval_progress.num_steps_completed
         )
         self.assertEqual(my_unit.steps_processed, steps_before_stopping)
 
@@ -129,9 +127,9 @@ def eval_step(
         state = init_eval_state(dataloader=dataloader)
         evaluate(state, my_unit)
 
-        self.assertEqual(state.eval_state.progress.num_epochs_completed, 1)
-        self.assertEqual(state.eval_state.progress.num_steps_completed_in_epoch, 0)
-        self.assertEqual(state.eval_state.progress.num_steps_completed, expected_steps)
+        self.assertEqual(my_unit.eval_progress.num_epochs_completed, 1)
+        self.assertEqual(my_unit.eval_progress.num_steps_completed_in_epoch, 0)
+        self.assertEqual(my_unit.eval_progress.num_steps_completed, expected_steps)
 
         # step_output should be reset to None
         self.assertEqual(state.eval_state.step_output, None)
@@ -148,7 +146,7 @@ def test_evaluate_with_callback(self) -> None:
         max_steps_per_epoch = 6
         expected_num_steps = dataset_len / batch_size
 
-        my_unit = MagicMock()
+        my_unit = DummyEvalUnit(2)
         dataloader = generate_random_dataloader(dataset_len, input_dim, batch_size)
         state = init_eval_state(
             dataloader=dataloader, max_steps_per_epoch=max_steps_per_epoch
@@ -225,7 +223,7 @@ def eval_step(
 
         assert state.eval_state
         if (
-            state.eval_state.progress.num_steps_completed_in_epoch + 1
+            self.eval_progress.num_steps_completed_in_epoch + 1
             == self.steps_before_stopping
         ):
             state.stop()