| """Train one model from a run configuration. |
| |
| python -m bifurcation.train --config records/euler/runs/cascaded_si_S10/config.yaml --data $EULER_DATA \\ |
| --out runs/euler_cascaded_si_S10 [--seed 0] [--device cuda] [--resume] |
| |
| The configuration names the dataset, the model, the method with its hyperparameters and the training schedule |
| (records/<system>/runs/<model>/config.yaml are the configurations of the paper's models). Before training, the |
| training pairs are written to OUT/manifest/: Cascaded SI with S segments reads the first 200 / S realizations of each |
| of the 80 training conditions, cut into S segments (16,000 pairs for every S; S = 1 is Direct SI); EDM reads the |
| endpoint pair (x_0, x_T) of all 200 realizations. Checkpoints are saved every ``training.val_every_updates`` updates |
| under OUT/checkpoints/ (bifurcation.trainer); ``--resume`` continues an interrupted run exactly from its |
| ``restart_latest.pt``. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import os |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| from .data import manifests as MF |
| from .data.loaders import get_dataloaders |
| from .registry import build_method, build_model |
| from .trainer import Trainer |
| from .utils import load_yaml, save_resolved_config, seed_everything |
|
|
| |
| RESUME_IDENTITY_KEYS = (("data", "name"), ("data", "packaged_dataset_path"), ("method", "name"), ("method", "data_mode"), |
| ("method", "sigma_train"), ("method", "P_mean"), ("method", "lr"), ("training", "batch_size"), |
| ("model", "name")) |
|
|
|
|
| def parse_args(argv=None): |
| p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| p.add_argument("--config", required=True, help="run configuration (YAML)") |
| p.add_argument("--data", required=True, help="dataset package directory") |
| p.add_argument("--out", required=True, help="run directory (new, or the one to resume)") |
| p.add_argument("--seed", type=int, default=0) |
| p.add_argument("--device", default=None, help="default: cuda if available, else cpu") |
| p.add_argument("--resume", action="store_true", help="continue from OUT/checkpoints/restart_latest.pt") |
| return p.parse_args(argv) |
|
|
|
|
| def write_training_manifest(cfg: dict, out: Path) -> Path: |
| """The training pairs of the run (see the module docstring); returns the manifest file.""" |
| packaged = cfg["data"]["packaged_dataset_path"] |
| n_ic, R = int(cfg["data"]["train_ics"]), int(cfg["data"]["train_realizations"]) |
| g = MF.fine_grid(packaged) |
| if len(g["gids_train"]) < n_ic or g["vecs_train_shape"][1] < R: |
| raise ValueError(f"the package holds {len(g['gids_train'])} training conditions x {g['vecs_train_shape'][1]} " |
| f"realizations, the configuration needs {n_ic} x {R}") |
| m = cfg["method"] |
| T = float(m["sampling"]["t_values"][-1]) |
| if m["data_mode"] == "fine_segment_pairs": |
| S = int(m["sampling"].get("n_segments", 1)) |
| if R % S: |
| raise ValueError(f"S = {S} must divide the {R} realizations per training condition") |
| geo = MF.segment_geometry(T, S, g["t_values"]) |
| rows = MF.block_rows_segment(g["gids_train"], n_ic, R // S) |
| name = f"segments_S{S:02d}" |
| path = MF.write_segment_manifest(out / "manifest", rows, geo, name, packaged, |
| extra={"realizations_per_condition": R // S}) |
| elif m["data_mode"] == "fine_endpoint_pairs": |
| geo = MF.segment_geometry(T, 1, g["t_values"]) |
| rows = MF.block_rows_endpoint(g["gids_train"], n_ic, R) |
| path = MF.write_endpoint_manifest(out / "manifest", rows, geo["n_T"], g["t_values"], "endpoints", packaged, |
| extra={"realizations_per_condition": R}) |
| else: |
| raise ValueError(f"unknown method.data_mode {m['data_mode']!r}") |
| if int(m["fine_horizon_idx"]) != int(geo["n_T"]): |
| raise ValueError(f"method.fine_horizon_idx {m['fine_horizon_idx']} != the index {geo['n_T']} of T = {T}") |
| return path |
|
|
|
|
| def main(argv=None) -> int: |
| args = parse_args(argv) |
| cfg = load_yaml(args.config) |
| cfg["data"]["packaged_dataset_path"] = str(Path(args.data).resolve()) + "/" |
| cfg["seed"] = int(args.seed) |
| run_dir = Path(args.out).resolve() |
| run_dir.mkdir(parents=True, exist_ok=True) |
| cfg["method"]["fine_slice_manifest"] = str(write_training_manifest(cfg, run_dir)) |
| cfg["training"]["device"] = str(args.device or os.environ.get("BIFURCATION_DEVICE") |
| or ("cuda" if torch.cuda.is_available() else "cpu")) |
|
|
| resume_payload = None |
| if args.resume: |
| restart_path = run_dir / "checkpoints" / "restart_latest.pt" |
| if not restart_path.is_file(): |
| raise FileNotFoundError(f"--resume: no restart checkpoint at {restart_path}") |
| resume_payload = torch.load(restart_path, map_location="cpu", weights_only=False) |
| old = resume_payload.get("config", {}) |
| for section, key in RESUME_IDENTITY_KEYS: |
| a, b = old.get(section, {}).get(key), cfg.get(section, {}).get(key) |
| if a is not None and a != b: |
| raise AssertionError(f"--resume: {section}.{key} was {a!r}, now {b!r}; refusing to continue a different run") |
| if int(old.get("seed", args.seed)) != int(args.seed): |
| raise AssertionError(f"--resume: the run has seed {old.get('seed')}, --seed {args.seed}") |
| print(f"resume: restart_latest.pt at step {resume_payload['step']}") |
| save_resolved_config(cfg, run_dir / "config_resolved.yaml") |
|
|
| seed_everything(int(args.seed)) |
| train_loader, val_loader = get_dataloaders(cfg, seed=int(args.seed)) |
| model = build_model(cfg["model"], cfg["method"], cfg["data"]) |
| method = build_method(cfg["method"]) |
| Trainer(model=model, method=method, train_loader=train_loader, val_loader=val_loader, cfg=cfg, run_dir=run_dir, |
| device=str(cfg["training"]["device"])).fit(resume=resume_payload) |
| print(f"\nrun: {run_dir}\ncheckpoints: {run_dir / 'checkpoints'}\nlog: {run_dir / 'train_log.csv'}") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
