class _Progress:
def __init__(self, show_limit: int = 50):
from rich.progress import Progress, BarColumn, TimeRemainingColumn
self._progress = Progress(
"[progress.description]{task.description}",
BarColumn(style="bar.back",
complete_style="bar.complete",
finished_style="bar.complete"),
"[progress.percentage]{task.percentage:>3.0f}%",
TimeRemainingColumn(),
auto_refresh=False)
self._last_update = _datetime.now()
self._show_limit = show_limit
self._completed = {}
self._have_entered = False
self._enter_args = None
def __enter__(self, *args, **kwargs):
self._enter_args = (args, kwargs)
return self
def __exit__(self, *args, **kwargs):
if self._have_entered:
self._progress.__exit__(*args, **kwargs)
return False
def add_task(self, description: str, total: int):
if total > self._show_limit:
if not self._have_entered:
args, kwargs = self._enter_args
self._progress.__enter__(*args, **kwargs)
self._have_entered = True
task_id = self._progress.add_task(description=description,
total=total)
self._completed[task_id] = 0
return task_id
else:
return None
def update(self, task_id: int, completed: float = None,
advance: float = None, force_update: bool = False):
if task_id is None:
return
elif completed is not None:
self._completed[task_id] = completed
elif advance is not None:
self._completed[task_id] += advance
else:
return
now = _datetime.now()
if force_update or (now - self._last_update).total_seconds() > 0.1:
self._progress.update(task_id=task_id,
completed=self._completed[task_id])
self._progress.refresh()
self._last_update = now
class tqdm_rich(std_tqdm): # pragma: no cover
"""Experimental rich.progress GUI version of tqdm!"""
# TODO: @classmethod: write()?
def __init__(self, *args, **kwargs):
"""
This class accepts the following parameters *in addition* to
the parameters accepted by `tqdm`.
Parameters
----------
progress : tuple, optional
arguments for `rich.progress.Progress()`.
"""
kwargs = kwargs.copy()
kwargs['gui'] = True
# convert disable = None to False
kwargs['disable'] = bool(kwargs.get('disable', False))
progress = kwargs.pop('progress', None)
super(tqdm_rich, self).__init__(*args, **kwargs)
if self.disable:
return
warn("rich is experimental/alpha",
TqdmExperimentalWarning,
stacklevel=2)
d = self.format_dict
if progress is None:
progress = ("[progress.description]{task.description}"
"[progress.percentage]{task.percentage:>4.0f}%",
BarColumn(bar_width=None),
FractionColumn(unit_scale=d['unit_scale'],
unit_divisor=d['unit_divisor']), "[",
TimeElapsedColumn(), "<", TimeRemainingColumn(), ",",
RateColumn(unit=d['unit'],
unit_scale=d['unit_scale'],
unit_divisor=d['unit_divisor']), "]")
self._prog = Progress(*progress, transient=not self.leave)
self._prog.__enter__()
self._task_id = self._prog.add_task(self.desc or "", **d)
def close(self, *args, **kwargs):
if self.disable:
return
super(tqdm_rich, self).close(*args, **kwargs)
self._prog.__exit__(None, None, None)
def clear(self, *_, **__):
pass
def display(self, *_, **__):
if not hasattr(self, '_prog'):
return
self._prog.update(self._task_id,
completed=self.n,
description=self.desc)
def reset(self, total=None):
"""
Resets to 0 iterations for repeated use.
Parameters
----------
total : int or float, optional. Total to use for the new bar.
"""
if hasattr(self, '_prog'):
self._prog.reset(total=total)
super(tqdm_rich, self).reset(total=total)
def __exit__(self, *args, **kwargs):
with self._lock:
for tid, task in self._tasks.items():
if not task.finished:
self.complete_task_nok(tid)
Progress.__exit__(self, *args, **kwargs)
class RichProgressBar(ProgressBarBase):
"""
Create a progress bar with `rich text formatting <https://github.com/willmcgugan/rich>`_.
Install it with pip:
.. code-block:: bash
pip install rich
.. code-block:: python
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import RichProgressBar
trainer = Trainer(callbacks=RichProgressBar())
Args:
refresh_rate: the number of updates per second, must be strictly positive
Raises:
ImportError:
If required `rich` package is not installed on the device.
"""
def __init__(self, refresh_rate: float = 1.0):
if not _RICH_AVAILABLE:
raise ImportError(
"`RichProgressBar` requires `rich` to be installed. Install it by running `pip install rich`."
)
super().__init__()
self._refresh_rate: float = refresh_rate
self._enabled: bool = True
self._total_val_batches: int = 0
self.progress: Progress = None
self.val_sanity_progress_bar_id: Optional[int] = None
self.main_progress_bar_id: Optional[int] = None
self.val_progress_bar_id: Optional[int] = None
self.test_progress_bar_id: Optional[int] = None
self.predict_progress_bar_id: Optional[int] = None
self.console = Console(record=True)
@property
def refresh_rate(self) -> int:
return self._refresh_rate
@property
def is_enabled(self) -> bool:
return self._enabled and self.refresh_rate > 0
@property
def is_disabled(self) -> bool:
return not self.is_enabled
def disable(self) -> None:
self._enabled = False
def enable(self) -> None:
self._enabled = True
@property
def sanity_check_description(self) -> str:
return "[Validation Sanity Check]"
@property
def validation_description(self) -> str:
return "[Validation]"
@property
def test_description(self) -> str:
return "[Testing]"
@property
def predict_description(self) -> str:
return "[Predicting]"
def setup(self, trainer, pl_module, stage):
self.progress = Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
BarColumn(),
BatchesProcessedColumn(),
"[",
CustomTimeColumn(),
ProcessingSpeedColumn(),
MetricsTextColumn(trainer, stage),
"]",
console=self.console,
refresh_per_second=self.refresh_rate,
).__enter__()
def on_sanity_check_start(self, trainer, pl_module):
super().on_sanity_check_start(trainer, pl_module)
self.val_sanity_progress_bar_id = self.progress.add_task(
f"[{STYLES['sanity_check']}]{self.sanity_check_description}",
total=trainer.num_sanity_val_steps,
)
def on_sanity_check_end(self, trainer, pl_module):
super().on_sanity_check_end(trainer, pl_module)
self.progress.update(self.val_sanity_progress_bar_id, visible=False)
def on_train_epoch_start(self, trainer, pl_module):
super().on_train_epoch_start(trainer, pl_module)
total_train_batches = self.total_train_batches
self._total_val_batches = self.total_val_batches
if total_train_batches != float("inf"):
# val can be checked multiple times per epoch
val_checks_per_epoch = total_train_batches // trainer.val_check_batch
self._total_val_batches = self._total_val_batches * val_checks_per_epoch
total_batches = total_train_batches + self._total_val_batches
train_description = self._get_train_description(trainer.current_epoch)
self.main_progress_bar_id = self.progress.add_task(
f"[{STYLES['train']}]{train_description}",
total=total_batches,
)
def on_validation_epoch_start(self, trainer, pl_module):
super().on_validation_epoch_start(trainer, pl_module)
if self._total_val_batches > 0:
self.val_progress_bar_id = self.progress.add_task(
f"[{STYLES['validate']}]{self.validation_description}",
total=self._total_val_batches,
)
def on_validation_epoch_end(self, trainer, pl_module):
super().on_validation_epoch_end(trainer, pl_module)
if self.val_progress_bar_id is not None:
self.progress.update(self.val_progress_bar_id, visible=False)
def on_test_epoch_start(self, trainer, pl_module):
super().on_train_epoch_start(trainer, pl_module)
self.test_progress_bar_id = self.progress.add_task(
f"[{STYLES['test']}]{self.test_description}",
total=self.total_test_batches,
)
def on_predict_epoch_start(self, trainer, pl_module):
super().on_predict_epoch_start(trainer, pl_module)
self.predict_progress_bar_id = self.progress.add_task(
f"[{STYLES['predict']}]{self.predict_description}",
total=self.total_predict_batches,
)
def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx,
dataloader_idx):
super().on_train_batch_end(trainer, pl_module, outputs, batch,
batch_idx, dataloader_idx)
if self._should_update(
self.train_batch_idx,
self.total_train_batches + self.total_val_batches):
self.progress.update(self.main_progress_bar_id, advance=1.0)
def on_validation_batch_end(self, trainer, pl_module, outputs, batch,
batch_idx, dataloader_idx):
super().on_validation_batch_end(trainer, pl_module, outputs, batch,
batch_idx, dataloader_idx)
if trainer.sanity_checking:
self.progress.update(self.val_sanity_progress_bar_id, advance=1.0)
elif self.val_progress_bar_id and self._should_update(
self.val_batch_idx,
self.total_train_batches + self.total_val_batches):
self.progress.update(self.main_progress_bar_id, advance=1.0)
self.progress.update(self.val_progress_bar_id, advance=1.0)
def on_test_batch_end(self, trainer, pl_module, outputs, batch, batch_idx,
dataloader_idx):
super().on_test_batch_end(trainer, pl_module, outputs, batch,
batch_idx, dataloader_idx)
if self._should_update(self.test_batch_idx, self.total_test_batches):
self.progress.update(self.test_progress_bar_id, advance=1.0)
def on_predict_batch_end(self, trainer, pl_module, outputs, batch,
batch_idx, dataloader_idx):
super().on_predict_batch_end(trainer, pl_module, outputs, batch,
batch_idx, dataloader_idx)
if self._should_update(self.predict_batch_idx,
self.total_predict_batches):
self.progress.update(self.predict_progress_bar_id, advance=1.0)
def _should_update(self, current, total) -> bool:
return self.is_enabled and (current % self.refresh_rate == 0
or current == total)
def _get_train_description(self, current_epoch: int) -> str:
train_description = f"[Epoch {current_epoch}]"
if len(self.validation_description) > len(train_description):
# Padding is required to avoid flickering due of uneven lengths of "Epoch X"
# and "Validation" Bar description
num_digits = len(str(current_epoch))
required_padding = (len(self.validation_description) -
len(train_description) + 1) - num_digits
for _ in range(required_padding):
train_description += " "
return train_description
def teardown(self, trainer, pl_module, stage):
self.progress.__exit__(None, None, None)
