"""Training logging and W&B naming hooks."""
import time
from pimm.utils.timer import Timer
from .default import HookBase
from .builder import HOOKS
LARGE_MODEL_OUTPUT_KEYS = {
"seg_logits",
"sem_logits",
"instance_embedding",
"vertex_embedding",
"sigma",
"point",
"pred_logits",
"pred_masks",
"teacher_logits",
}
def model_output_scalar_keys(model_output_dict):
"""Return model-output keys that InformationWriter should log."""
has_total_loss = "total_loss" in model_output_dict
keys = [
key
for key in model_output_dict.keys()
if (
key not in LARGE_MODEL_OUTPUT_KEYS
and "match_" not in key
and not (has_total_loss and key == "loss")
and key != "total_loss"
)
]
if has_total_loss:
keys.append("loss")
return keys
def model_output_scalar_value(model_output_dict, key):
"""Return a scalar float for a selected model-output key, or None."""
val = model_output_dict["total_loss"] if key == "loss" and "total_loss" in model_output_dict else model_output_dict[key]
try:
if hasattr(val, "item") and callable(getattr(val, "item", None)):
return float(val.item())
if isinstance(val, (int, float)):
return float(val)
return float(val)
except Exception:
return None
[docs]
@HOOKS.register_module()
class WandbNamer(HookBase):
"""Auto-generate ``cfg.wandb_run_name`` from selected config values.
Joins the resolved values of the requested (possibly nested) config keys
with a separator, optionally formatting large numbers with K/M/B suffixes,
then appends any ``extra`` strings, to build a descriptive W&B run name.
Runs once in ``modify_config`` (before the writer is built), so the name is
in place before logging starts. Registered as ``WandbNamer``.
Args:
keys (tuple): Config keys to include in the name, in order. Supports
dotted nested access such as ``"data.train.max_len"`` or
``"model.type"``; keys that resolve to ``None`` are skipped.
Defaults to ``()``.
sep (str): Separator joining the parts. Defaults to ``"-"``.
format_numbers (bool): If ``True``, render numeric values with K/M/B
suffixes (e.g. ``1000000`` -> ``"1M"``). Defaults to ``True``.
extra (str or tuple, optional): Extra string(s) appended after the
keyed parts, e.g. ``"fft"`` or ``"scratch"``. A bare string is
wrapped to a one-tuple. Defaults to ``None``.
Note:
If ``wandb_run_name`` was set on the command line (present in
``cfg._cli_options``), this hook leaves it untouched. So
``--options wandb_run_name=my-name`` overrides the auto-generated name.
Example:
Add to ``cfg.hooks``; once in ``modify_config`` (before the writer is
built) it sets ``cfg.wandb_run_name`` from the joined config values:
.. code-block:: python
hooks = [
dict(type="WandbNamer",
keys=("model.type", "data.train.max_len", "seed"),
extra="fft"),
]
# → sets cfg.wandb_run_name = "Sonata-v1m1-1M-0-fft" (left untouched if
# wandb_run_name was passed on the CLI)
"""
def __init__(self, keys=(), sep="-", format_numbers=True, extra=None):
"""Store config keys and formatting options for generated run names."""
self.keys = keys
self.sep = sep
self.format_numbers = format_numbers
# normalize extra to tuple
if extra is None:
self.extra = ()
elif isinstance(extra, str):
self.extra = (extra,)
else:
self.extra = tuple(extra)
def _get_nested(self, cfg, key_path):
"""Get nested config value: 'data.train.max_len' -> cfg.data['train']['max_len']"""
parts = key_path.split('.')
val = cfg
for part in parts:
if hasattr(val, part):
val = getattr(val, part)
elif isinstance(val, dict) and part in val:
val = val[part]
else:
return None
return val
def _format_value(self, val):
"""Format a value for the run name."""
if self.format_numbers and isinstance(val, (int, float)):
return self._format_number(val)
return str(val)
def _format_number(self, n):
"""Format large numbers with suffixes (K, M, B)."""
n_val = float(n)
if abs(n_val) >= 1_000_000_000:
return f"{n_val / 1_000_000_000:.1f}B".rstrip('0').rstrip('.')
elif abs(n_val) >= 1_000_000:
return f"{n_val / 1_000_000:.1f}M".rstrip('0').rstrip('.')
elif abs(n_val) >= 1_000:
return f"{n_val / 1_000:.1f}K".rstrip('0').rstrip('.')
return str(int(n_val) if n_val == int(n_val) else n_val)
[docs]
def modify_config(self, cfg):
"""Build wandb_run_name from specified keys."""
# skip if wandb_run_name already set via CLI
if 'wandb_run_name' in getattr(cfg, '_cli_options', set()):
return
parts = []
for key in self.keys:
val = self._get_nested(cfg, key)
if val is not None:
parts.append(self._format_value(val))
# append extra strings
parts.extend(self.extra)
if parts:
cfg.wandb_run_name = self.sep.join(parts)
[docs]
@HOOKS.register_module()
class IterationTimer(HookBase):
"""Measure data/batch latency and append timing and ETA to iteration logs.
Times each training iteration and feeds the results into
``trainer.storage`` and the per-iteration log string. In ``before_step`` it
records ``data_time`` (time spent waiting for the batch); in ``after_step``
it records ``batch_time`` (full iteration), decrements the remaining-
iteration counter, and appends ``Data``, ``Batch``, and ``Remain`` (ETA,
computed from the running average batch time) to ``comm_info["iter_info"]``.
The remaining-iteration budget is initialized in ``before_train`` and the
per-iteration timer is reset in ``before_epoch``. Registered as
``IterationTimer``.
Args:
warmup_iter (int): Number of initial iterations whose timing is reset
(excluded from the running averages) so startup/compile overhead
does not skew ETA. Defaults to ``1``.
Example:
Add to ``cfg.hooks``; it times each iteration and feeds the timing into
``trainer.storage`` and the per-step log line:
.. code-block:: python
hooks = [dict(type="IterationTimer", warmup_iter=2)]
# → puts scalars "data_time" (before_step) and "batch_time" (after_step)
# into trainer.storage and appends "Data … Batch … Remain HH:MM:SS"
# (ETA) to comm_info["iter_info"]; the first 2 iters reset the averages
"""
def __init__(self, warmup_iter=1):
"""Configure how many initial iterations are excluded from averages."""
self._warmup_iter = warmup_iter
self._start_time = time.perf_counter()
self._iter_timer = Timer()
self._remain_iter = 0
[docs]
def before_train(self):
"""Initialize remaining-iteration accounting at train start."""
self._start_time = time.perf_counter()
_remain_epoch = self.trainer.max_epoch - self.trainer.start_epoch
self._remain_iter = _remain_epoch * len(self.trainer.train_loader)
[docs]
def before_epoch(self):
"""Reset the per-iteration timer at the start of each epoch."""
self._iter_timer.reset()
[docs]
def before_step(self):
"""Record data loading time before the model consumes the batch."""
data_time = self._iter_timer.seconds()
self.trainer.storage.put_scalar("data_time", data_time)
[docs]
def after_step(self):
"""Record batch time and append timing fields to iter_info."""
batch_time = self._iter_timer.seconds()
self._iter_timer.reset()
self.trainer.storage.put_scalar("batch_time", batch_time)
self._remain_iter -= 1
remain_time = self._remain_iter * self.trainer.storage.history("batch_time").avg
t_m, t_s = divmod(remain_time, 60)
t_h, t_m = divmod(t_m, 60)
remain_time = "{:02d}:{:02d}:{:02d}".format(int(t_h), int(t_m), int(t_s))
if "iter_info" in self.trainer.comm_info.keys():
info = (
"Data {data_time_val:.3f} ({data_time_avg:.3f}) "
"Batch {batch_time_val:.3f} ({batch_time_avg:.3f}) "
"Remain {remain_time} ".format(
data_time_val=self.trainer.storage.history("data_time").val,
data_time_avg=self.trainer.storage.history("data_time").avg,
batch_time_val=self.trainer.storage.history("batch_time").val,
batch_time_avg=self.trainer.storage.history("batch_time").avg,
remain_time=remain_time,
)
)
self.trainer.comm_info["iter_info"] += info
if self.trainer.comm_info["iter"] <= self._warmup_iter:
self.trainer.storage.history("data_time").reset()
self.trainer.storage.history("batch_time").reset()
