File size: 4,038 Bytes
d5147ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | 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])))
# Dedup by step (keep last)
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, [])
# Replace same-step point if re-logged (eval retries etc.)
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
# Stable order
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()
|