File size: 3,881 Bytes
45507bb | 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 | """Lightweight experiment tracking: one directory per run under `runs/`.
runs/<timestamp>-<name>/
config.yaml resolved config
env.json git commit, hardware, library versions, runtime profile
metrics.csv step-level training/validation metrics
tb/ TensorBoard events (if tensorboard is installed)
summary.json final metrics, timings, cost
checkpoints/ last.pt (resumable), best.pt, final.pt (calibrated)
"""
from __future__ import annotations
import csv
import json
import platform
import subprocess
import time
from pathlib import Path
import yaml
def git_commit() -> str | None:
try:
out = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, text=True, timeout=5,
cwd=Path(__file__).resolve().parent)
dirty = subprocess.run(["git", "status", "--porcelain"], capture_output=True, text=True, timeout=5,
cwd=Path(__file__).resolve().parent).stdout.strip()
return (out.stdout.strip() + ("-dirty" if dirty else "")) if out.returncode == 0 else None
except Exception:
return None
def environment_info() -> dict:
import numpy
info = {"python": platform.python_version(), "platform": platform.platform(),
"processor": platform.processor(), "numpy": numpy.__version__, "git_commit": git_commit()}
try:
import torch
info["torch"] = torch.__version__
info["cuda"] = torch.version.cuda
if torch.cuda.is_available():
info["gpu"] = torch.cuda.get_device_name(0)
except ImportError:
pass
return info
class RunLogger:
def __init__(self, run_dir: str | Path, cfg: dict | None = None, resume: bool = False):
self.dir = Path(run_dir)
self.dir.mkdir(parents=True, exist_ok=True)
(self.dir / "checkpoints").mkdir(exist_ok=True)
if cfg is not None and not (resume and (self.dir / "config.yaml").exists()):
(self.dir / "config.yaml").write_text(yaml.safe_dump(cfg, sort_keys=False))
self._csv_path = self.dir / "metrics.csv"
self._fields: list[str] | None = None
if resume and self._csv_path.exists():
with open(self._csv_path) as f:
self._fields = next(csv.reader(f), None)
self.tb = None
try:
from torch.utils.tensorboard import SummaryWriter
self.tb = SummaryWriter(str(self.dir / "tb"))
except Exception:
pass
def write_json(self, name: str, obj) -> None:
(self.dir / name).write_text(json.dumps(obj, indent=2, default=str))
def log(self, step: int, **metrics) -> None:
row = {"step": step, "time": round(time.time(), 1), **metrics}
if self._fields is None:
self._fields = list(row)
with open(self._csv_path, "w", newline="") as f:
csv.writer(f).writerow(self._fields)
for k in row:
if k not in self._fields: # new metric: rewrite header
self._fields.append(k)
rows = list(csv.DictReader(open(self._csv_path)))
with open(self._csv_path, "w", newline="") as f:
w = csv.DictWriter(f, self._fields)
w.writeheader()
w.writerows(rows)
with open(self._csv_path, "a", newline="") as f:
csv.DictWriter(f, self._fields).writerow(row)
if self.tb is not None:
for k, v in metrics.items():
if isinstance(v, (int, float)):
self.tb.add_scalar(k, v, step)
def close(self) -> None:
if self.tb is not None:
self.tb.close()
def new_run_dir(name: str, root: str | Path = "runs") -> Path:
return Path(root) / f"{time.strftime('%Y%m%d-%H%M%S')}-{name}"
|