Dedup training pipeline, add seeding/resume and per-epoch metrics logging
cli.py and scripts/train.py duplicated ~140 lines of training setup and had drifted (scripts/train.py forgot to save model_config, breaking predict on those checkpoints). Extract shared logic into giant/constants.py (X_DIM, target names), giant/config.py (device/git/TOML/seeding helpers, run metadata), and giant/pipeline.py (the actual training-job orchestration), so both entry points become thin CLI wrappers around the same code path. Also adds --seed/--resume support (checkpoints now carry optimizer/scheduler state, epoch, and best_val_loss), a richer [meta] section in the saved config.toml (git hash, seed, versions, timestamp, invocation, dataset stats), and a metrics.csv (train/val loss, lr, epoch time) written every epoch and append-safe across resumes. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
+12
-160
@@ -1,54 +1,10 @@
|
||||
import argparse
|
||||
import subprocess
|
||||
import tomllib
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from giant.data.loader import (
|
||||
find_parquet_files,
|
||||
load_event_ids,
|
||||
iter_file_chunks,
|
||||
build_index_maps_from_files,
|
||||
)
|
||||
from giant.data.transforms import build_features, _WelfordAccumulator
|
||||
from giant.data.dataset import make_event_split, StreamingStepsDataset
|
||||
from giant.model.network import DenoisingMLP
|
||||
from giant.train import train as run_training
|
||||
|
||||
|
||||
def _auto_device() -> torch.device:
|
||||
if torch.cuda.is_available():
|
||||
return torch.device("cuda")
|
||||
if torch.backends.mps.is_available():
|
||||
return torch.device("mps")
|
||||
return torch.device("cpu")
|
||||
|
||||
|
||||
def _git_hash() -> str:
|
||||
try:
|
||||
return subprocess.check_output(
|
||||
["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL
|
||||
).decode().strip()
|
||||
except Exception:
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _load_config(path: str) -> dict:
|
||||
with open(path, "rb") as f:
|
||||
return tomllib.load(f)
|
||||
|
||||
|
||||
def _save_config(cfg: dict, out_dir: Path) -> None:
|
||||
lines = [f"# git: {_git_hash()}", ""]
|
||||
for section, values in cfg.items():
|
||||
lines.append(f"[{section}]")
|
||||
for k, v in values.items():
|
||||
lines.append(f"{k:<12} = {repr(v) if isinstance(v, str) else v}")
|
||||
lines.append("")
|
||||
(out_dir / "config.toml").write_text("\n".join(lines))
|
||||
from giant import config as gconfig
|
||||
from giant.pipeline import run_train_job
|
||||
|
||||
|
||||
def main() -> None:
|
||||
@@ -63,42 +19,28 @@ def main() -> None:
|
||||
parser.add_argument("--n-blocks", type=int)
|
||||
parser.add_argument("--emb-dim", type=int)
|
||||
parser.add_argument("--val-fraction", type=float)
|
||||
parser.add_argument("--seed", type=int, help="Random seed for reproducibility")
|
||||
parser.add_argument("--shuffle-buffer", type=int, default=65536,
|
||||
help="Rows held in RAM for shuffling per worker (default: 65536)")
|
||||
parser.add_argument("--out", default=None, help="Checkpoint output directory (default: auto from hyperparams)")
|
||||
parser.add_argument("--device", default=None, help="cpu | cuda | mps (default: auto)")
|
||||
parser.add_argument("--num-workers", type=int)
|
||||
parser.add_argument("--resume", default=None, help="Checkpoint .pt to resume training from")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Defaults → config file → explicit CLI flags.
|
||||
cfg: dict = {
|
||||
"train": {
|
||||
"mode": "flow", "epochs": 100, "batch_size": 4096, "lr": 3e-4,
|
||||
"val_fraction": 0.1, "num_workers": 4,
|
||||
},
|
||||
"model": {
|
||||
"hidden_dim": 256, "n_blocks": 6, "emb_dim": 16,
|
||||
},
|
||||
}
|
||||
|
||||
if args.config:
|
||||
file_cfg = _load_config(args.config)
|
||||
for section in ("train", "model"):
|
||||
cfg[section].update(file_cfg.get(section, {}))
|
||||
|
||||
cli_train = {k: v for k, v in {
|
||||
"mode": args.mode, "epochs": args.epochs, "batch_size": args.batch_size,
|
||||
"lr": args.lr, "val_fraction": args.val_fraction, "num_workers": args.num_workers,
|
||||
"seed": args.seed,
|
||||
}.items() if v is not None}
|
||||
cli_model = {k: v for k, v in {
|
||||
"hidden_dim": args.hidden_dim, "n_blocks": args.n_blocks, "emb_dim": args.emb_dim,
|
||||
}.items() if v is not None}
|
||||
cfg["train"].update(cli_train)
|
||||
cfg["model"].update(cli_model)
|
||||
|
||||
config_path = Path(args.config) if args.config else None
|
||||
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config_path, cli_train, cli_model)
|
||||
t, m = cfg["train"], cfg["model"]
|
||||
|
||||
device = torch.device(args.device) if args.device else _auto_device()
|
||||
device = torch.device(args.device) if args.device else gconfig.auto_device()
|
||||
out_dir = Path(args.out or (
|
||||
f"checkpoints/{t['mode']}"
|
||||
f"_h{m['hidden_dim']}"
|
||||
@@ -111,100 +53,10 @@ def main() -> None:
|
||||
print(f"device: {device}")
|
||||
print(f"out_dir: {out_dir}")
|
||||
|
||||
# --- Discover files ---
|
||||
files = find_parquet_files(args.data)
|
||||
print(f"found {len(files)} parquet file(s)")
|
||||
|
||||
# --- Scan event IDs (single column, cheap) ---
|
||||
print("scanning event IDs …")
|
||||
all_event_ids = np.concatenate([load_event_ids(f) for f in files])
|
||||
train_events, val_events = make_event_split(all_event_ids, val_fraction=t["val_fraction"])
|
||||
n_train_steps = (np.isin(all_event_ids, np.array(sorted(train_events)))).sum()
|
||||
print(f" {len(all_event_ids):,} steps | "
|
||||
f"{len(train_events)} train events (~{n_train_steps:,} steps) | "
|
||||
f"{len(val_events)} val events")
|
||||
|
||||
# --- Scan PDG / material vocabularies (2 columns, cheap) ---
|
||||
print("building vocabulary maps …")
|
||||
pdg_map, mat_map = build_index_maps_from_files(files)
|
||||
print(f" {len(pdg_map)} PDG codes | {len(mat_map)} materials")
|
||||
|
||||
# --- Streaming normalizer fit over training data ---
|
||||
print("fitting normalizer (streaming) …")
|
||||
events_arr = np.array(sorted(train_events))
|
||||
cond_acc = _WelfordAccumulator(9)
|
||||
tgt_acc = _WelfordAccumulator(9)
|
||||
for path in files:
|
||||
for chunk in iter_file_chunks(path):
|
||||
mask = np.isin(chunk["event_id"], events_arr)
|
||||
if not mask.any():
|
||||
continue
|
||||
chunk_tr = {k: v[mask] for k, v in chunk.items()}
|
||||
cond_cont, _, target, _, _ = build_features(chunk_tr, pdg_map, mat_map)
|
||||
cond_acc.update(cond_cont)
|
||||
tgt_acc.update(target)
|
||||
cond_norm = cond_acc.to_normalizer()
|
||||
tgt_norm = tgt_acc.to_normalizer()
|
||||
|
||||
# --- Streaming datasets ---
|
||||
train_ds = StreamingStepsDataset(
|
||||
files=files,
|
||||
split_events=train_events,
|
||||
pdg_map=pdg_map,
|
||||
mat_map=mat_map,
|
||||
cond_normalizer=cond_norm,
|
||||
target_normalizer=tgt_norm,
|
||||
batch_size=t["batch_size"],
|
||||
shuffle_buffer=args.shuffle_buffer,
|
||||
shuffle=True,
|
||||
)
|
||||
val_ds = StreamingStepsDataset(
|
||||
files=files,
|
||||
split_events=val_events,
|
||||
pdg_map=pdg_map,
|
||||
mat_map=mat_map,
|
||||
cond_normalizer=cond_norm,
|
||||
target_normalizer=tgt_norm,
|
||||
batch_size=t["batch_size"],
|
||||
shuffle=False,
|
||||
)
|
||||
|
||||
# Dataset yields whole batches already, so batch_size=None tells DataLoader
|
||||
# to pass them through instead of re-collating row-by-row in Python.
|
||||
pin = device.type == "cuda"
|
||||
train_loader = DataLoader(
|
||||
train_ds, batch_size=None,
|
||||
num_workers=t["num_workers"], pin_memory=pin,
|
||||
)
|
||||
val_loader = DataLoader(
|
||||
val_ds, batch_size=None,
|
||||
num_workers=t["num_workers"], pin_memory=pin,
|
||||
)
|
||||
|
||||
model = DenoisingMLP(
|
||||
pdg_vocab=len(pdg_map),
|
||||
mat_vocab=len(mat_map),
|
||||
hidden_dim=m["hidden_dim"],
|
||||
n_blocks=m["n_blocks"],
|
||||
emb_dim=m["emb_dim"],
|
||||
)
|
||||
print(f"model: {sum(p.numel() for p in model.parameters()):,} parameters")
|
||||
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
_save_config(cfg, out_dir)
|
||||
|
||||
run_training(
|
||||
model=model,
|
||||
train_loader=train_loader,
|
||||
val_loader=val_loader,
|
||||
mode=t["mode"],
|
||||
epochs=t["epochs"],
|
||||
lr=t["lr"],
|
||||
device=device,
|
||||
out_dir=out_dir,
|
||||
normalizer_dict={"cond": cond_norm.to_dict(), "target": tgt_norm.to_dict()},
|
||||
pdg_map={str(k): v for k, v in pdg_map.items()},
|
||||
mat_map={str(k): v for k, v in mat_map.items()},
|
||||
run_train_job(
|
||||
data=Path(args.data), cfg=cfg, out_dir=out_dir, device=device,
|
||||
shuffle_buffer=args.shuffle_buffer, num_workers=t["num_workers"],
|
||||
resume=Path(args.resume) if args.resume else None, echo=print,
|
||||
)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user