File size: 3,423 Bytes
a10f8a4 5b9021b a10f8a4 d6bd312 a10f8a4 d6bd312 a10f8a4 d605416 a10f8a4 d605416 a10f8a4 d6bd312 a10f8a4 | 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 | """Checkpoint save/load for resumable OmniAL training."""
import json
import os
from pathlib import Path
from typing import Any, Dict, Optional
import torch
from torch.amp import GradScaler
def ensure_run_dir(run_dir: str) -> Path:
p = Path(run_dir)
(p / "checkpoints").mkdir(parents=True, exist_ok=True)
(p / "metrics").mkdir(parents=True, exist_ok=True)
return p
def save_checkpoint(
path: str,
epoch: int,
model: torch.nn.Module,
optimizer: torch.optim.Optimizer,
scaler: Optional[GradScaler],
history: Dict[str, Any],
best_metric: float,
cfg: dict,
tag: str = "",
scheduler=None,
) -> str:
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
payload = {
"epoch": epoch,
"model_state": model.state_dict(),
"optimizer_state": optimizer.state_dict(),
"scaler_state": scaler.state_dict() if scaler is not None else None,
"scheduler_state": scheduler.state_dict() if scheduler is not None else None,
"history": history,
"best_metric": best_metric,
"cfg_train": cfg.get("train", {}),
"tag": tag,
}
torch.save(payload, path)
return path
def load_checkpoint(
path: str,
model: torch.nn.Module,
optimizer: Optional[torch.optim.Optimizer] = None,
scaler: Optional[GradScaler] = None,
device: str = "cpu",
) -> Dict[str, Any]:
# Full training checkpoints store history/optimizer (numpy scalars etc.).
# PyTorch 2.6+ defaults weights_only=True, which rejects those — use False
# only for checkpoints you trust (your own training runs).
ckpt = torch.load(path, map_location=device, weights_only=False)
model.load_state_dict(ckpt["model_state"])
if optimizer is not None and ckpt.get("optimizer_state"):
optimizer.load_state_dict(ckpt["optimizer_state"])
if scaler is not None and ckpt.get("scaler_state"):
try:
scaler.load_state_dict(ckpt["scaler_state"])
except Exception as e:
print(f"[WARN] could not restore GradScaler state: {e}")
return ckpt
def restore_scheduler(scheduler, ckpt: Dict[str, Any]) -> None:
if scheduler is None or not ckpt.get("scheduler_state"):
return
try:
scheduler.load_state_dict(ckpt["scheduler_state"])
except Exception as e:
print(f"[WARN] could not restore scheduler state: {e}")
def save_metrics_json(run_dir: str, history: Dict[str, Any], filename: str = "metrics_history.json") -> str:
path = Path(run_dir) / "metrics" / filename
with open(path, "w", encoding="utf-8") as f:
json.dump(history, f, indent=2, ensure_ascii=False)
return str(path)
def append_epoch_metrics(run_dir: str, epoch_record: Dict[str, Any]) -> str:
metrics_dir = Path(run_dir) / "metrics"
metrics_dir.mkdir(parents=True, exist_ok=True)
path = metrics_dir / "epoch_metrics.jsonl"
with open(path, "a", encoding="utf-8") as f:
f.write(json.dumps(epoch_record, ensure_ascii=False) + "\n")
return str(path)
def find_latest_checkpoint(run_dir: str) -> Optional[str]:
ckpt_dir = Path(run_dir) / "checkpoints"
if not ckpt_dir.is_dir():
return None
last = ckpt_dir / "last.pt"
if last.is_file():
return str(last)
numbered = sorted(ckpt_dir.glob("epoch_*.pt"), key=lambda p: p.stat().st_mtime)
return str(numbered[-1]) if numbered else None
|