File size: 1,946 Bytes
3495881 | 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 | """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
|