73 lines
3.3 KiB
Python
73 lines
3.3 KiB
Python
# Pause-aware tqdm progress bar shared by the training scripts.
|
|
"""tqdm subclass whose elapsed time and rate skip explicitly excluded windows."""
|
|
from contextlib import contextmanager
|
|
|
|
from tqdm.auto import tqdm
|
|
|
|
__all__ = ["PauseAwareTqdm"]
|
|
|
|
|
|
class PauseAwareTqdm(tqdm):
|
|
r"""
|
|
tqdm whose elapsed time and rate ignore the time spent inside `paused()` blocks.
|
|
|
|
tqdm has no notion of a pause: the smoothed step rate is derived from the
|
|
time since `last_print_t` and the closing line re-derives itself from the
|
|
overall average since `start_t`. A non-training pause (checkpoint save,
|
|
validation) left uncompensated lands in the interval of the next update, so
|
|
the bar shows that step as much slower, and the pause dilutes the closing
|
|
line. Every tqdm read goes through the instance clock `self._time` (an
|
|
instance attribute holding `time.time` by default), so subtracting the
|
|
total pause time there keeps the elapsed time and every rate reporting the
|
|
training speed only.
|
|
|
|
`close()` additionally rebases `start_t` onto the smoothed rate before the
|
|
final display: the closing line is re-derived as the overall average, which
|
|
one-off costs (worker warm-up, first dataloader fetch) would otherwise
|
|
dilute even when all pauses are excluded.
|
|
|
|
Verified against tqdm 4.27 - 4.67: `tqdm.auto`, which the training scripts
|
|
import, first ships in 4.27, and the pause compensation works from there.
|
|
The `close()` rebase additionally needs `format_dict` (4.30+); below that
|
|
it is skipped while elapsed time and rates are still compensated. If a
|
|
future release drops an internal this relies on, the class degrades to a
|
|
plain bar instead of raising: the probe below leaves `_base_time` unset
|
|
and `paused()` / the `close()` rebase become no-ops.
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
base = getattr(self, "_time", None) # tqdm's own clock (time.time by default)
|
|
self._base_time = base if callable(base) else None
|
|
self._pause_total = 0.0
|
|
# `initial` only became an instance attribute in tqdm ~4.55; keep a private
|
|
# copy so the close() rebase below works on older releases as well.
|
|
initial = getattr(self, "initial", None)
|
|
self._initial = kwargs.get("initial", 0) if initial is None else initial
|
|
if self._base_time is not None:
|
|
def clock():
|
|
return self._base_time() - self._pause_total
|
|
self._time = clock
|
|
|
|
@contextmanager
|
|
def paused(self):
|
|
"""Keeps the wall time of the block out of the elapsed time and the rates."""
|
|
if self.disable or self._base_time is None:
|
|
yield
|
|
return
|
|
started = self._base_time()
|
|
try:
|
|
yield
|
|
finally:
|
|
self._pause_total += self._base_time() - started
|
|
|
|
def close(self):
|
|
# tqdm re-derives the closing line as the overall average (elapsed from bar
|
|
# creation); rebasing start_t onto the smoothed rate the bar was just showing
|
|
# makes that line repeat the pace the run actually sustained.
|
|
if not self.disable and self._base_time is not None:
|
|
rate = getattr(self, "format_dict", {}).get("rate")
|
|
if rate:
|
|
self.start_t = self._time() - (self.n - self._initial) / rate
|
|
super().close()
|