Add streaming data pipeline and giant CLI entry point
- Streaming pipeline: row-group-level parquet reading (PyArrow) so large files never fully land in RAM; Welford online algorithm for normalizer fitting; StreamingStepsDataset with shuffle buffer and multi-worker file striping; event-ID scan and vocab scan via cheap single-column reads - giant/cli.py: typer-based CLI with `giant train` subcommand, mirroring scripts/train.py; --shuffle-buffer flag for RAM control - pyproject.toml: add typer>=0.12 dependency and giant entry point - train.py: replace len(loader.dataset) with local counters (compatible with IterableDataset which has no __len__) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
+202
@@ -0,0 +1,202 @@
|
||||
import subprocess
|
||||
import tomllib
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import typer
|
||||
from torch.utils.data import DataLoader
|
||||
from typing_extensions import Annotated
|
||||
|
||||
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
|
||||
|
||||
app = typer.Typer(no_args_is_help=True)
|
||||
|
||||
|
||||
@app.callback()
|
||||
def _main() -> None:
|
||||
"""GIANT — Geant4 step-function surrogate."""
|
||||
|
||||
|
||||
class Mode(str, Enum):
|
||||
flow = "flow"
|
||||
ddpm = "ddpm"
|
||||
|
||||
|
||||
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_toml(path: Path) -> 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))
|
||||
|
||||
|
||||
@app.command()
|
||||
def train(
|
||||
data: Annotated[Path, typer.Argument(help="Parquet file or directory of parquet files")],
|
||||
config: Annotated[Optional[Path], typer.Option(help="TOML config file (overridden by explicit flags)")] = None,
|
||||
mode: Annotated[Optional[Mode], typer.Option(help="Generative model: flow matching or DDPM")] = None,
|
||||
epochs: Annotated[Optional[int], typer.Option()] = None,
|
||||
batch_size: Annotated[Optional[int], typer.Option()] = None,
|
||||
lr: Annotated[Optional[float], typer.Option()] = None,
|
||||
hidden_dim: Annotated[Optional[int], typer.Option()] = None,
|
||||
n_blocks: Annotated[Optional[int], typer.Option()] = None,
|
||||
emb_dim: Annotated[Optional[int], typer.Option()] = None,
|
||||
val_fraction: Annotated[Optional[float], typer.Option()] = None,
|
||||
shuffle_buffer: Annotated[int, typer.Option(help="Rows held in RAM per worker for shuffling")] = 65536,
|
||||
out: Annotated[Optional[Path], typer.Option(help="Checkpoint dir (default: auto from hyperparams)")] = None,
|
||||
device: Annotated[Optional[str], typer.Option(help="cpu | cuda | mps (default: auto)")] = None,
|
||||
num_workers: Annotated[Optional[int], typer.Option()] = None,
|
||||
) -> None:
|
||||
"""Train the GIANT surrogate model."""
|
||||
# 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 config is not None:
|
||||
file_cfg = _load_toml(config)
|
||||
for section in ("train", "model"):
|
||||
cfg[section].update(file_cfg.get(section, {}))
|
||||
|
||||
cli_train = {k: v for k, v in {
|
||||
"mode": mode.value if mode is not None else None,
|
||||
"epochs": epochs, "batch_size": batch_size, "lr": lr,
|
||||
"val_fraction": val_fraction, "num_workers": num_workers,
|
||||
}.items() if v is not None}
|
||||
cli_model = {k: v for k, v in {
|
||||
"hidden_dim": hidden_dim, "n_blocks": n_blocks, "emb_dim": emb_dim,
|
||||
}.items() if v is not None}
|
||||
cfg["train"].update(cli_train)
|
||||
cfg["model"].update(cli_model)
|
||||
|
||||
t, m = cfg["train"], cfg["model"]
|
||||
|
||||
_device = torch.device(device) if device else _auto_device()
|
||||
out_dir = out or Path(
|
||||
f"checkpoints/{t['mode']}"
|
||||
f"_h{m['hidden_dim']}"
|
||||
f"_b{m['n_blocks']}"
|
||||
f"_e{m['emb_dim']}"
|
||||
f"_lr{t['lr']}"
|
||||
f"_bs{t['batch_size']}"
|
||||
)
|
||||
|
||||
typer.echo(f"device: {_device}")
|
||||
typer.echo(f"out_dir: {out_dir}")
|
||||
|
||||
files = find_parquet_files(data)
|
||||
typer.echo(f"found {len(files)} parquet file(s)")
|
||||
|
||||
typer.echo("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"])
|
||||
events_arr = np.array(sorted(train_events))
|
||||
n_train_steps = int(np.isin(all_event_ids, events_arr).sum())
|
||||
typer.echo(
|
||||
f" {len(all_event_ids):,} steps | "
|
||||
f"{len(train_events)} train events (~{n_train_steps:,} steps) | "
|
||||
f"{len(val_events)} val events"
|
||||
)
|
||||
|
||||
typer.echo("building vocabulary maps …")
|
||||
pdg_map, mat_map = build_index_maps_from_files(files)
|
||||
typer.echo(f" {len(pdg_map)} PDG codes | {len(mat_map)} materials")
|
||||
|
||||
typer.echo("fitting normalizer (streaming) …")
|
||||
cond_acc = _WelfordAccumulator(9)
|
||||
tgt_acc = _WelfordAccumulator(6)
|
||||
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()
|
||||
|
||||
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,
|
||||
shuffle_buffer=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,
|
||||
shuffle=False,
|
||||
)
|
||||
|
||||
pin = _device.type == "cuda"
|
||||
train_loader = DataLoader(
|
||||
train_ds, batch_size=t["batch_size"],
|
||||
num_workers=t["num_workers"], pin_memory=pin,
|
||||
)
|
||||
val_loader = DataLoader(
|
||||
val_ds, batch_size=t["batch_size"],
|
||||
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"],
|
||||
)
|
||||
typer.echo(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()},
|
||||
)
|
||||
Reference in New Issue
Block a user