train_monitor.py

"""
TrainMonitor — лёгкий мониторинг лосса и стабильности градиентов.
Встраивается в тренировочный цикл 3 вызовами:
    monitor = TrainMonitor(log_every=25)
    ...
    monitor.step(loss, model, global_step)   # внутри цикла
    ...
    monitor.summary()                        # после тренировки
"""

import time
import numpy as np
import torch

# ---------- floating-point friendly helpers ----------
def _f(val):
    """округление до 6 знаков, безопасное к None."""
    if val is None:
        return None
    return round(float(val), 6)


class TrainMonitor:
    def __init__(
        self,
        log_every: int = 50,
        ema_alpha: float = 0.05,
        window: int = 100,
        warmup_steps: int = 5,
        csv_path: str | None = "monitor_log.csv",
    ):
        self.log_every = log_every
        self.ema_alpha = ema_alpha
        self.window = window
        self.warmup_steps = warmup_steps
        self.csv_path = csv_path

        # loss
        self.loss_ema = None
        self.loss_start = None
        self.loss_history = []         # (step, raw_loss)
        self.loss_ema_history = []     # (step, ema)

        # grad
        self.grad_norm_history = []    # (step, norm)
        self.grad_mean_history = []    # (step, mean_abs)
        self.grad_std_history = []     # (step, std)
        self.grad_dead_pct_history = []# (step, dead_%)

        # timing
        self.step_times = []
        self.last_time = None

        # CSV header
        if self.csv_path:
            with open(self.csv_path, "w") as f:
                f.write("step,loss,loss_ema,grad_norm,grad_mean,grad_std,grad_dead_pct,loss_delta_pct,step_ms\n")

    # ------------------------------------------------------------------
    # вызывается КАЖДЫЙ шаг (на main process)
    # ------------------------------------------------------------------
    def step(self, loss: torch.Tensor, model: torch.nn.Module, step: int):
        loss_val = loss.detach().item()

        # --- timing ---
        now = time.perf_counter()
        if self.last_time is not None:
            self.step_times.append((step, (now - self.last_time) * 1000))  # ms
        self.last_time = now

        # --- EMA loss ---
        if self.loss_ema is None:
            self.loss_ema = loss_val
        else:
            self.loss_ema = self.loss_ema * (1 - self.ema_alpha) + loss_val * self.ema_alpha

        if self.loss_start is None and step >= self.warmup_steps:
            self.loss_start = self.loss_ema

        self.loss_history.append((step, loss_val))
        self.loss_ema_history.append((step, self.loss_ema))

        # --- gradient statistics (только когда есть градиенты) ---
        grad_norm = 0.0
        grad_mean = 0.0
        grad_std = 0.0
        grad_dead_pct = 0.0

        total_elems = 0
        total_abs = 0.0
        total_sq = 0.0
        dead_elems = 0
        grad_norm_sq = 0.0

        has_grads = False
        for p in model.parameters():
            if p.grad is None:
                continue
            has_grads = True
            g = p.grad.detach().float()
            n = g.numel()
            total_elems += n
            grad_norm_sq += g.pow(2).sum().item()
            total_abs += g.abs().sum().item()
            total_sq += g.pow(2).sum().item()
            dead_elems += (g.abs() < 1e-8).sum().item()

        if has_grads and total_elems > 0:
            grad_norm = np.sqrt(grad_norm_sq)
            grad_mean = total_abs / total_elems
            # std = sqrt(E[g^2] - E[|g|]^2) — приближение через mean_abs, корректнее через mean(g^2)
            mean_sq = total_sq / total_elems
            grad_std = np.sqrt(max(0, mean_sq - grad_mean ** 2))
            grad_dead_pct = (dead_elems / total_elems) * 100.0

        self.grad_norm_history.append((step, grad_norm))
        self.grad_mean_history.append((step, grad_mean))
        self.grad_std_history.append((step, grad_std))
        self.grad_dead_pct_history.append((step, grad_dead_pct))

        # --- печать и CSV ---
        if step % self.log_every == 0:
            self._log(step)

        # --- CSV ---
        if self.csv_path and has_grads:
            loss_delta = 0.0
            if self.loss_start is not None and self.loss_start > 0:
                loss_delta = ((self.loss_start - self.loss_ema) / self.loss_start) * 100.0
            step_ms = self.step_times[-1][1] if self.step_times else 0.0
            with open(self.csv_path, "a") as f:
                f.write(
                    f"{step},{_f(loss_val)},{_f(self.loss_ema)},"
                    f"{_f(grad_norm)},{_f(grad_mean)},{_f(grad_std)},{_f(grad_dead_pct)},"
                    f"{_f(loss_delta)},{_f(step_ms)}\n"
                )

    # ------------------------------------------------------------------
    # внутренняя печать
    # ------------------------------------------------------------------
    def _log(self, step: int):
        loss_delta = 0.0
        if self.loss_start is not None and self.loss_start > 0:
            loss_delta = ((self.loss_start - self.loss_ema) / self.loss_start) * 100.0

        # последние значения градиентов
        gn, gm, gs, gd = 0.0, 0.0, 0.0, 0.0
        if self.grad_norm_history:
            gn = self.grad_norm_history[-1][1]
            gm = self.grad_mean_history[-1][1]
            gs = self.grad_std_history[-1][1]
            gd = self.grad_dead_pct_history[-1][1]

        # среднее время шага за последние log_every шагов
        recent_times = [t for s, t in self.step_times[-self.log_every:]]
        avg_ms = np.mean(recent_times) if recent_times else 0.0

        # rolling loss delta (последние window шагов)
        roll_delta = 0.0
        if len(self.loss_ema_history) >= self.window:
            old_ema = self.loss_ema_history[-self.window][1]
            if old_ema > 0:
                roll_delta = ((old_ema - self.loss_ema) / old_ema) * 100.0

        arrow = "↓" if loss_delta > 0 else ("↑" if loss_delta < 0 else "→")
        rarrow = "↓" if roll_delta > 0 else ("↑" if roll_delta < 0 else "→")

        print(
            f"[step {step:>6d}] "
            f"loss={self.loss_ema:.6f} "
            f"(raw={self.loss_history[-1][1]:.6f}) "
            f"{arrow}{abs(loss_delta):.2f}% "
            f"{rarrow}{abs(roll_delta):.2f}%/{self.window} "
            f"| grad: norm={gn:.3f} μ={gm:.6f} σ={gs:.6f} dead={gd:.1f}% "
            f"| {avg_ms:.0f}ms/step"
        )

    # ------------------------------------------------------------------
    # конец эпохи — средний лосс, градиенты
    # ------------------------------------------------------------------
    def end_epoch(self, epoch: int, step: int):
        if not self.loss_history:
            return
        # собираем метрики за последние N шагов (вся эпоха — от последнего end_epoch или от начала)
        # для простоты: последние self.window шагов, но не больше чем вся история
        n = min(self.window, len(self.loss_history))
        recent_losses = [v for _, v in self.loss_history[-n:]]
        recent_emas  = [v for _, v in self.loss_ema_history[-n:]]
        recent_gn    = [v for _, v in self.grad_norm_history[-n:]]
        recent_dead  = [v for _, v in self.grad_dead_pct_history[-n:]]

        avg_loss = np.mean(recent_losses)
        avg_ema  = np.mean(recent_emas)
        avg_gn   = np.mean(recent_gn)
        avg_dead = np.mean(recent_dead)

        loss_delta = 0.0
        if self.loss_start is not None and self.loss_start > 0:
            loss_delta = ((self.loss_start - self.loss_ema) / self.loss_start) * 100.0

        arrow = "↓" if loss_delta > 0 else ("↑" if loss_delta < 0 else "→")
        print(
            f"\n{'═' * 60}\n"
            f"  ЭПОХА {epoch + 1} завершена (шаг {step})\n"
            f"  Средний loss (raw):    {avg_loss:.6f}\n"
            f"  Средний loss (EMA):    {avg_ema:.6f}\n"
            f"  Текущий loss (EMA):    {self.loss_ema:.6f}\n"
            f"  Падение от старта:     {arrow}{abs(loss_delta):.2f}%\n"
            f"  Средний grad_norm:     {avg_gn:.4f}\n"
            f"  Средний dead grads:    {avg_dead:.1f}%\n"
            f"{'═' * 60}\n"
        )

    # ------------------------------------------------------------------
    # итоговый отчёт
    # ------------------------------------------------------------------
    def summary(self):
        if not self.loss_history:
            print("[Monitor] нет данных.")
            return

        steps = [s for s, _ in self.loss_history]
        losses = [v for _, v in self.loss_history]
        ema_losses = [v for _, v in self.loss_ema_history]
        grad_norms = [v for _, v in self.grad_norm_history]

        total_delta = 0.0
        if self.loss_start is not None and self.loss_start > 0:
            total_delta = ((self.loss_start - self.loss_ema) / self.loss_start) * 100.0

        arrow = "↓" if total_delta > 0 else ("↑" if total_delta < 0 else "→")
        print("\n" + "=" * 72)
        print("  TRAIN MONITOR — ИТОГОВЫЙ ОТЧЁТ")
        print("=" * 72)
        print(f"  Всего шагов:           {len(steps)}")
        print(f"  Loss start (EMA):      {self.loss_start:.6f}" if self.loss_start else "  Loss start:            (warmup...)")
        print(f"  Loss final (EMA):      {self.loss_ema:.6f}")
        print(f"  Loss final (raw):      {losses[-1]:.6f}")
        print(f"  Падение лосса:         {arrow}{abs(total_delta):.2f}%")
        print()
        print(f"  Grad norm (среднее):   {np.mean(grad_norms):.4f}")
        print(f"  Grad norm (std):       {np.std(grad_norms):.4f}")
        print(f"  Grad norm (max):       {np.max(grad_norms):.4f}")
        print(f"  Grad norm (min):       {np.min(grad_norms):.4f}")
        print(f"  Grad norm / median:    {np.median(grad_norms):.4f}")
        print()
        deads = [v for _, v in self.grad_dead_pct_history]
        if deads:
            print(f"  Dead grads (среднее):  {np.mean(deads):.2f}%")
            print(f"  Dead grads (max):      {np.max(deads):.2f}%")
        print()
        if self.step_times:
            times = [t for _, t in self.step_times]
            print(f"  Step time (среднее):   {np.mean(times):.0f} ms")
            print(f"  Step time (p99):       {np.percentile(times, 99):.0f} ms")
        print()
        print(f"  Стабильность градиентов:")
        if len(grad_norms) > 1:
            cv = np.std(grad_norms) / (np.mean(grad_norms) + 1e-8)  # coefficient of variation
            verdict = "✅ отлично" if cv < 0.3 else ("⚠️ умеренно" if cv < 0.6 else "❌ нестабильно")
            print(f"    CV grad_norm:         {cv:.3f}  {verdict}")
        loss_std = np.std(ema_losses[-self.window:]) if len(ema_losses) >= self.window else np.std(ema_losses)
        loss_mean = np.mean(ema_losses[-self.window:]) if len(ema_losses) >= self.window else np.mean(ema_losses)
        if len(ema_losses) > 1:
            loss_cv = loss_std / (loss_mean + 1e-8)
            verdict = "✅ отлично" if loss_cv < 0.02 else ("⚠️ умеренно" if loss_cv < 0.05 else "❌ нестабильно")
            print(f"    CV loss (окно):       {loss_cv:.4f}  {verdict}")
        print("=" * 72 + "\n")

        if self.csv_path:
            print(f"[Monitor] CSV сохранён: {self.csv_path}")

    # ------------------------------------------------------------------
    # для внешнего использования: получить текущие метрики словарём
    # ------------------------------------------------------------------
    def get_metrics(self) -> dict:
        out = {"loss_ema": self.loss_ema, "step": self.loss_history[-1][0] if self.loss_history else 0}
        if self.grad_norm_history:
            out["grad_norm"] = self.grad_norm_history[-1][1]
            out["grad_mean"] = self.grad_mean_history[-1][1]
            out["grad_std"] = self.grad_std_history[-1][1]
            out["grad_dead_pct"] = self.grad_dead_pct_history[-1][1]
        return out