github-actions[bot]
60 lines
2.6 KiB
Python
60 lines
2.6 KiB
Python
"""Sigma-bucketed loss tracking.
|
|
Maps each training step's per-element sigmas and losses to buckets.
|
|
Smoothing is left to wandb's UI.
|
|
"""
|
|
|
|
import bisect
|
|
from collections import defaultdict
|
|
|
|
|
|
class SigmaBucketTracker:
|
|
"""Map per-element sigma values to named buckets for per-bucket loss logging.
|
|
By default, partitions [0, 1] into four equal-width buckets.
|
|
Custom boundaries can be provided for non-uniform bucketing.
|
|
Each call to update() receives per-element sigmas and losses (both [B,]),
|
|
buckets each element, and computes the mean loss per bucket. This gives
|
|
accurate per-sigma loss tracking even for batch_size > 1.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
bucket_boundaries: list[float] | None = None,
|
|
) -> None:
|
|
if bucket_boundaries is None:
|
|
bucket_boundaries = [0.0, 0.25, 0.5, 0.75, 1.0]
|
|
if len(bucket_boundaries) < 2:
|
|
raise ValueError("bucket_boundaries must have at least 2 elements")
|
|
if any(bucket_boundaries[i] >= bucket_boundaries[i + 1] for i in range(len(bucket_boundaries) - 1)):
|
|
raise ValueError("bucket_boundaries must be strictly increasing")
|
|
self._boundaries = list(bucket_boundaries)
|
|
self._num_buckets = len(bucket_boundaries) - 1
|
|
self._bucket_labels = [
|
|
f"{bucket_boundaries[i]:.2f}-{bucket_boundaries[i + 1]:.2f}" for i in range(self._num_buckets)
|
|
]
|
|
self._last_metrics: dict[str, float] = {}
|
|
|
|
def _get_bucket_index(self, sigma: float) -> int:
|
|
"""Map sigma value to bucket index."""
|
|
idx = bisect.bisect_right(self._boundaries, sigma) - 1
|
|
return max(0, min(idx, self._num_buckets - 1))
|
|
|
|
def update(self, sigmas: list[float], losses: list[float]) -> None:
|
|
"""Record per-element losses into their sigma buckets.
|
|
Args:
|
|
sigmas: Per-element sigma values, one per batch element.
|
|
losses: Per-element losses, one per batch element.
|
|
"""
|
|
if not sigmas:
|
|
self._last_metrics = {}
|
|
return
|
|
bucket_losses: dict[int, list[float]] = defaultdict(list)
|
|
for sigma, loss in zip(sigmas, losses, strict=True):
|
|
bucket_losses[self._get_bucket_index(sigma)].append(loss)
|
|
self._last_metrics = {self._bucket_labels[b]: sum(vals) / len(vals) for b, vals in bucket_losses.items()}
|
|
|
|
def get_metrics(self, prefix: str = "train") -> dict[str, float]:
|
|
"""Return the mean loss for each bucket hit on the last update.
|
|
Wandb handles smoothing in the UI.
|
|
"""
|
|
return {f"{prefix}/loss_sigma_{label}": loss for label, loss in self._last_metrics.items()}
|