add warmup steps to iteration time logger (#837)

Summary:
Pull Request resolved: #837

Adding warmup steps to different performance loggers

Reviewed By: diego-urgell

Differential Revision: D57595957

fbshipit-source-id: e92bf0345ed93ea944cb751bdfd162535c69c332

Author: galrotem

tests/framework/callbacks/test_iteration_time_logger.py

@@ -27,6 +27,7 @@
 from torchtnt.framework.state import EntryPoint, PhaseState, State
 from torchtnt.framework.train import _train_impl, train
 from torchtnt.utils.loggers.logger import MetricLogger
+from torchtnt.utils.timer import Timer
 
 
 class IterationTimeLoggerTest(unittest.TestCase):
@@ -164,3 +165,34 @@ def test_with_summary_writer(self) -> None:
         train(my_unit, dataloader, max_epochs=2, callbacks=[callback])
         # 2 epochs, 6 iterations each, logging every third step
         self.assertEqual(logger.add_scalar.call_count, 4)
+
+    def test_warmup_steps(self) -> None:
+        logger = MagicMock(spec=MetricLogger)
+        callback = IterationTimeLogger(logger=logger, warmup_steps=1)
+        timer = Timer()
+        timer.recorded_durations = {"train_iteration_time": [1, 2]}
+
+        # ensure that we don't log for the first step
+        callback._log_step_metrics("train_iteration_time", timer, 1)
+        logger.log.assert_not_called()
+
+        # second step should log
+        callback._log_step_metrics("train_iteration_time", timer, 2)
+        self.assertEqual(logger.log.call_count, 1)
+
+    def test_invalid_params(self) -> None:
+        logger = MagicMock(spec=MetricLogger)
+        with self.assertRaisesRegex(
+            ValueError, "moving_avg_window must be at least 1, got 0"
+        ):
+            IterationTimeLogger(logger=logger, moving_avg_window=0)
+
+        with self.assertRaisesRegex(
+            ValueError, "log_every_n_steps must be at least 1, got -1"
+        ):
+            IterationTimeLogger(logger=logger, log_every_n_steps=-1)
+
+        with self.assertRaisesRegex(
+            ValueError, "warmup_steps must be at least 0, got -1"
+        ):
+            IterationTimeLogger(logger=logger, warmup_steps=-1)

torchtnt/framework/callbacks/iteration_time_logger.py

@@ -27,20 +27,37 @@ class IterationTimeLogger(Callback):
     Args:
         logger: Either a :class:`torchtnt.loggers.tensorboard.TensorBoardLogger`
             or a :class:`torch.utils.tensorboard.SummaryWriter` instance.
-        moving_avg_window: an optional int to control the moving average window
-        log_every_n_steps: an optional int to control the log frequency
+        moving_avg_window: an int to control the moving average window. Default is 1.
+        log_every_n_steps: an int to control the log frequency. Default is 1.
+        warmup_steps: an int to control the number of warmup steps. We will start logging only after the amount of warmup steps were completed. Default is 0.
     """
 
     def __init__(
         self,
         logger: Union[MetricLogger, SummaryWriter],
+        *,
         moving_avg_window: int = 1,
         log_every_n_steps: int = 1,
+        warmup_steps: int = 0,
     ) -> None:
         self._logger = logger
+
+        if moving_avg_window < 1:
+            raise ValueError(
+                f"moving_avg_window must be at least 1, got {moving_avg_window}"
+            )
         self.moving_avg_window = moving_avg_window
+
+        if log_every_n_steps < 1:
+            raise ValueError(
+                f"log_every_n_steps must be at least 1, got {log_every_n_steps}"
+            )
         self.log_every_n_steps = log_every_n_steps
 
+        if warmup_steps < 0:
+            raise ValueError(f"warmup_steps must be at least 0, got {warmup_steps}")
+        self.warmup_steps = warmup_steps
+
     @rank_zero_fn
     def _log_step_metrics(
         self,
@@ -53,6 +70,9 @@ def _log_step_metrics(
         was configured.
 
         """
+        if step_logging_for <= self.warmup_steps:
+            return
+
         if step_logging_for % self.log_every_n_steps != 0:
             return