| from __future__ import annotations |
|
|
| import json |
| import os |
| from typing import Any |
|
|
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import torch |
| from torch.utils.tensorboard import SummaryWriter |
|
|
|
|
| class TensorBoardLogger: |
| """TensorBoard + PNG exporter with resume-safe history. |
| |
| Scalars are persisted to ``{plot_dir}/../tensorboard_scalars.json`` so that |
| pause→resume training redraws the full curve instead of only the current segment. |
| """ |
|
|
| def __init__(self, log_dir: str, plot_dir: str): |
| self.writer = SummaryWriter(log_dir=log_dir) |
| self.plot_dir = plot_dir |
| self.scalars: dict[str, list[tuple[int, float]]] = {} |
| self._history_path = os.path.join(os.path.dirname(plot_dir.rstrip("/")), "tensorboard_scalars.json") |
| self._load_history() |
|
|
| def _load_history(self) -> None: |
| if not os.path.isfile(self._history_path): |
| return |
| try: |
| with open(self._history_path, encoding="utf-8") as f: |
| raw = json.load(f) |
| for tag, points in raw.items(): |
| cleaned: list[tuple[int, float]] = [] |
| for item in points: |
| if not isinstance(item, (list, tuple)) or len(item) != 2: |
| continue |
| cleaned.append((int(item[0]), float(item[1]))) |
| |
| by_step = {s: v for s, v in cleaned} |
| self.scalars[tag] = sorted(by_step.items(), key=lambda x: x[0]) |
| print( |
| f"[TB] Restored {sum(len(v) for v in self.scalars.values())} scalar points " |
| f"from {self._history_path}", |
| flush=True, |
| ) |
| except Exception as exc: |
| print(f"[TB] WARN: failed to load history {self._history_path}: {exc}", flush=True) |
|
|
| def _save_history(self) -> None: |
| os.makedirs(os.path.dirname(self._history_path) or ".", exist_ok=True) |
| payload = {tag: [[s, v] for s, v in points] for tag, points in self.scalars.items()} |
| tmp = self._history_path + ".tmp" |
| with open(tmp, "w", encoding="utf-8") as f: |
| json.dump(payload, f) |
| os.replace(tmp, self._history_path) |
|
|
| @staticmethod |
| def _to_scalar(value: Any) -> float | None: |
| if isinstance(value, bool): |
| return float(value) |
| if isinstance(value, (int, float)): |
| return float(value) |
| if isinstance(value, torch.Tensor): |
| if value.numel() != 1: |
| return None |
| return float(value.detach().cpu().item()) |
| return None |
|
|
| def log(self, metrics: dict[str, Any], step: int) -> None: |
| for name, value in metrics.items(): |
| scalar = self._to_scalar(value) |
| if scalar is None: |
| continue |
| self.writer.add_scalar(name, scalar, step) |
| pts = self.scalars.setdefault(name, []) |
| |
| if pts and pts[-1][0] == step: |
| pts[-1] = (step, scalar) |
| else: |
| pts.append((step, scalar)) |
|
|
| def save_plots(self) -> None: |
| if not self.scalars: |
| return |
|
|
| self._save_history() |
| os.makedirs(self.plot_dir, exist_ok=True) |
|
|
| for tag, points in self.scalars.items(): |
| if not points: |
| continue |
| |
| points = sorted(points, key=lambda x: x[0]) |
| steps, values = zip(*points) |
| fig, ax = plt.subplots(figsize=(8, 4)) |
| ax.plot(steps, values, linewidth=1.5) |
| ax.set_xlabel("step") |
| ax.set_ylabel(tag) |
| ax.set_title(tag) |
| ax.grid(True, alpha=0.3) |
| fig.tight_layout() |
| out = os.path.join(self.plot_dir, f"{tag.replace('/', '_')}.png") |
| fig.savefig(out) |
| plt.close(fig) |
|
|
| def close(self) -> None: |
| self.save_plots() |
| self.writer.close() |
|
|
