| """Checkpointing: resumable experiments, never rely on RAM.""" |
| from __future__ import annotations |
|
|
| import json |
| import os |
| import time |
| from dataclasses import asdict |
|
|
| import torch |
|
|
| from .model import TinyChess, TinyChessConfig |
|
|
|
|
| def save_checkpoint(path, model, cfg, optimizer=None, scheduler=None, step=0, |
| metrics=None, extra=None, rng=True): |
| tmp = path + ".tmp" |
| payload = { |
| "model": model.state_dict(), |
| "cfg": cfg.to_dict() if hasattr(cfg, "to_dict") else dict(cfg), |
| "step": step, |
| "metrics": metrics or {}, |
| "extra": extra or {}, |
| "param_report": model.param_report(), |
| "saved_at": time.time(), |
| } |
| if optimizer is not None: |
| payload["optimizer"] = optimizer.state_dict() |
| if scheduler is not None: |
| payload["scheduler"] = scheduler.state_dict() |
| if rng: |
| payload["rng"] = { |
| "torch": torch.get_rng_state(), |
| "python": __import__("random").getstate(), |
| "numpy": __import__("numpy").random.get_state(), |
| } |
| os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) |
| torch.save(payload, tmp) |
| os.replace(tmp, path) |
| return path |
|
|
|
|
| def load_checkpoint(path, map_location="cpu", restore_rng=True): |
| ck = torch.load(path, map_location=map_location, weights_only=False) |
| cfg = TinyChessConfig(**ck["cfg"]) |
| model = TinyChess(cfg) |
| model.load_state_dict(ck["model"]) |
| if restore_rng and "rng" in ck: |
| torch.set_rng_state(ck["rng"]["torch"]) |
| __import__("random").setstate(ck["rng"]["python"]) |
| __import__("numpy").random.set_state(ck["rng"]["numpy"]) |
| return model, cfg, ck |
|
|
|
|
| def load_model_only(path, map_location="cpu"): |
| ck = torch.load(path, map_location=map_location, weights_only=False) |
| cfg = TinyChessConfig(**ck["cfg"]) |
| model = TinyChess(cfg) |
| model.load_state_dict(ck["model"]) |
| model.eval() |
| return model, cfg |
|
|
