File size: 6,285 Bytes
41dd7e0 | 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 104 105 106 107 108 109 110 111 112 113 114 115 116 | """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
# a restart may only continue a run with the same data, method, model and seed
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())
|