| """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]: |
| |
| |
| |
| 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 |
|
|
