Allow setting mininterval for TQDM progress bar (meta-pytorch#987)

Summary:
Pull Request resolved: meta-pytorch#987

This uses `tqdm`'s exposes built-in `mininterval` to the callback -- this allows the user to control the TQDM update rate based on the number of seconds rather than the number of iteration steps.

Reviewed By: JKSenthil

Differential Revision: D72068311

fbshipit-source-id: 91105da3db1dbadac5b89f9865406f75035dccb6

Author: alanhdu

torchtnt/framework/callbacks/tqdm_progress_bar.py

@@ -30,15 +30,19 @@ class TQDMProgressBar(Callback):
 
     Args:
         refresh_rate: Determines at which rate (in number of steps) the progress bars get updated.
+        mininterval: Minimum display update interval (in seconds). If None, use TQDM's default.
         file: specifies where to output the progress messages (default: sys.stderr)
     """
 
     def __init__(
         self,
         refresh_rate: int = 1,
         file: Optional[Union[TextIO, io.StringIO]] = None,
+        *,
+        mininterval: float | None = None,
     ) -> None:
         self._refresh_rate = refresh_rate
+        self._mininterval = mininterval
         self._file = file
 
         self._train_progress_bar: Optional[tqdm] = None
@@ -56,6 +60,7 @@ def on_train_epoch_start(self, state: State, unit: TTrainUnit) -> None:
                 max_steps=train_state.max_steps,
                 max_steps_per_epoch=train_state.max_steps_per_epoch,
                 file=self._file,
+                mininterval=self._mininterval,
             )
 
     def on_train_step_end(self, state: State, unit: TTrainUnit) -> None:
@@ -87,6 +92,7 @@ def on_eval_epoch_start(self, state: State, unit: TEvalUnit) -> None:
                 max_steps=eval_state.max_steps,
                 max_steps_per_epoch=eval_state.max_steps_per_epoch,
                 file=self._file,
+                mininterval=self._mininterval,
             )
 
     def on_eval_step_end(self, state: State, unit: TEvalUnit) -> None:

torchtnt/utils/tqdm.py

@@ -25,6 +25,7 @@ def create_progress_bar(
     num_steps_completed: int,
     max_steps: Optional[int],
     max_steps_per_epoch: Optional[int],
+    mininterval: float | None = None,
     file: Optional[Union[TextIO, io.StringIO]] = None,
 ) -> tqdm:
     """Constructs a :func:`tqdm` progress bar. The number of steps in an epoch is inferred from the dataloader, num_steps_completed, max_steps and max_steps_per_epoch.
@@ -36,6 +37,7 @@ def create_progress_bar(
         num_steps_completed: an integer for the number of steps completed so far in the loop.
         max_steps: an optional integer for the number of max steps in the loop.
         max_steps_per_epoch: an optional integer for the number of max steps per epoch.
+        mininterval: Minimum display update interval (in seconds). If None, use TQDM's default.
         file: specifies where to output the progress messages (default: sys.stderr)
     """
     current_epoch = num_epochs_completed
@@ -45,12 +47,16 @@ def create_progress_bar(
         max_steps=max_steps,
         max_steps_per_epoch=max_steps_per_epoch,
     )
+    kwargs = {}
+    if mininterval is not None:
+        kwargs["mininterval"] = mininterval
     return tqdm(
         desc=f"{desc} {current_epoch}",
         total=total,
         initial=num_steps_completed,
         bar_format="{l_bar}{bar}{r_bar}\n",
         file=file,
+        **kwargs,
     )