From 74d0883868420887917993e7ed54ae0a92f5de50 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Fri, 19 Jun 2026 13:17:15 +0200 Subject: [PATCH] Remove scripts/train.py in favor of the giant train CLI The Typer-based giant/cli.py train command now has full feature parity (dropout, warmup-epochs, validate-steps, shorthand flags), making the standalone argparse script redundant. --- CLAUDE.md | 4 +- README.md | 6 --- giant/cli.py | 90 +++++++++++++++++++++++++++---------- scripts/train.py | 114 ----------------------------------------------- 4 files changed, 69 insertions(+), 145 deletions(-) delete mode 100644 scripts/train.py diff --git a/CLAUDE.md b/CLAUDE.md index f7a1a8c..71e91fa 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -9,8 +9,8 @@ uv sync --extra cpu # install dependencies with CPU-only torch (s uv sync --extra cuda # install dependencies with CUDA 11.8 torch uv sync --extra cpu --extra dev # add dev extras (pytest, etc.) pytest # run tests -python scripts/train.py --data path/to/steps.parquet --mode flow # train (flow matching) -python scripts/train.py --data path/to/steps.parquet --mode ddpm # train (DDPM baseline) +giant train path/to/steps.parquet --mode flow # train (flow matching) +giant train path/to/steps.parquet --mode ddpm # train (DDPM baseline) ``` `cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x; newer torch requires newer NVIDIA drivers). Plain `uv sync` with no extra will not install torch at all; uv has no concept of a "default extra", so `--extra cpu` should always be included unless you need GPU support. diff --git a/README.md b/README.md index 8059ad3..e0b3264 100644 --- a/README.md +++ b/README.md @@ -72,12 +72,6 @@ uv sync --extra cpu --extra dev # add dev tools (pytest, ruff, ty) ## Training -```bash -python scripts/train.py --data path/to/steps.parquet --mode flow -``` - -or via the installed CLI: - ```bash giant train path/to/steps.parquet --mode flow giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt diff --git a/giant/cli.py b/giant/cli.py index 389c3d0..d71e640 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -59,45 +59,73 @@ def train( ], config: Annotated[ Optional[Path], - typer.Option(help="TOML config file (overridden by explicit flags)"), + typer.Option( + "--config", "-c", help="TOML config file (overridden by explicit flags)" + ), ] = None, mode: Annotated[ - Optional[Mode], typer.Option(help="Generative model: flow matching or DDPM") + Optional[Mode], + typer.Option("--mode", "-m", help="Generative model: flow matching or DDPM"), + ] = None, + epochs: Annotated[Optional[int], typer.Option("--epochs", "-e")] = None, + batch_size: Annotated[Optional[int], typer.Option("--batch-size", "-b")] = None, + lr: Annotated[Optional[float], typer.Option("--lr", "-l")] = None, + warmup_epochs: Annotated[ + Optional[int], typer.Option("--warmup-epochs", "-w") + ] = None, + hidden_dim: Annotated[Optional[int], typer.Option("--hidden-dim", "-H")] = None, + n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None, + emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None, + dropout: Annotated[ + Optional[float], + typer.Option( + "--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)" + ), + ] = None, + val_fraction: Annotated[ + Optional[float], typer.Option("--val-fraction", "-f") ] = 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, seed: Annotated[ - Optional[int], typer.Option(help="Random seed for reproducibility") + Optional[int], + typer.Option("--seed", "-s", help="Random seed for reproducibility"), ] = None, validate_every: Annotated[ Optional[int], - typer.Option(help="Run marginal+KL validation every N epochs (0 disables)"), + typer.Option( + "--validate-every", + "-v", + help="Run marginal+KL validation every N epochs (0 disables)", + ), ] = None, validate_steps: Annotated[ Optional[int], typer.Option( + "--validate-steps", + "-t", help="Flow matching ODE steps used during marginal validation " - "(ignored in ddpm mode, which always runs the full schedule)" + "(ignored in ddpm mode, which always runs the full schedule)", ), ] = None, shuffle_buffer: Annotated[ - int, typer.Option(help="Rows held in RAM per worker for shuffling") + int, + typer.Option( + "--shuffle-buffer", "-B", help="Rows held in RAM per worker for shuffling" + ), ] = 65536, out: Annotated[ Optional[Path], - typer.Option(help="Checkpoint dir (default: auto from hyperparams)"), + typer.Option( + "--out", "-o", help="Checkpoint dir (default: auto from hyperparams)" + ), ] = None, device: Annotated[ - Optional[str], typer.Option(help="cpu | cuda | mps (default: auto)") + Optional[str], + typer.Option("--device", "-D", help="cpu | cuda | mps (default: auto)"), ] = None, - num_workers: Annotated[Optional[int], typer.Option()] = None, + num_workers: Annotated[Optional[int], typer.Option("--num-workers", "-j")] = None, resume: Annotated[ - Optional[Path], typer.Option(help="Checkpoint .pt to resume training from") + Optional[Path], + typer.Option("--resume", "-r", help="Checkpoint .pt to resume training from"), ] = None, ) -> None: """Train the GIANT surrogate model.""" @@ -108,6 +136,7 @@ def train( "epochs": epochs, "batch_size": batch_size, "lr": lr, + "warmup_epochs": warmup_epochs, "val_fraction": val_fraction, "num_workers": num_workers, "seed": seed, @@ -122,6 +151,7 @@ def train( "hidden_dim": hidden_dim, "n_blocks": n_blocks, "emb_dim": emb_dim, + "dropout": dropout, }.items() if v is not None } @@ -161,27 +191,41 @@ def predict( Path, typer.Argument(help="Parquet file or directory of parquet files") ], checkpoint: Annotated[ - Path, typer.Option(help="Path to checkpoint .pt file (best.pt or last.pt)") + Path, + typer.Option( + "--checkpoint", + "-c", + help="Path to checkpoint .pt file (best.pt or last.pt)", + ), ], coord: Annotated[ Coord, typer.Option( + "--coord", + "-C", help="global: full physical units, world frame (default). " "local: raw 9D model output (denormalised only, local frame, " "log-scaled scalars) alongside the matching ground-truth target " - "for the same input file — requires post-step columns." + "for the same input file — requires post-step columns.", ), ] = Coord.global_, out: Annotated[ Optional[Path], typer.Option( - help="Output parquet path (default: _predicted[_local].parquet)" + "--out", + "-o", + help="Output parquet path (default: _predicted[_local].parquet)", ), ] = None, - batch_size: Annotated[int, typer.Option(help="Inference batch size")] = 4096, - steps: Annotated[int, typer.Option(help="Flow matching ODE steps")] = 10, + batch_size: Annotated[ + int, typer.Option("--batch-size", "-b", help="Inference batch size") + ] = 4096, + steps: Annotated[ + int, typer.Option("--steps", "-s", help="Flow matching ODE steps") + ] = 10, device: Annotated[ - Optional[str], typer.Option(help="cpu | cuda | mps (default: auto)") + Optional[str], + typer.Option("--device", "-d", help="cpu | cuda | mps (default: auto)"), ] = None, ) -> None: """Run trained model on a parquet file and save predictions.""" diff --git a/scripts/train.py b/scripts/train.py deleted file mode 100644 index 899f380..0000000 --- a/scripts/train.py +++ /dev/null @@ -1,114 +0,0 @@ -import argparse -from pathlib import Path - -import torch - -from giant import config as gconfig -from giant.pipeline import run_train_job - - -def main() -> None: - parser = argparse.ArgumentParser(description="Train GIANT surrogate model") - parser.add_argument("--config", default=None, help="Path to TOML config file") - parser.add_argument( - "--data", required=True, help="Path to parquet file or directory" - ) - parser.add_argument("--mode", choices=["flow", "ddpm"]) - parser.add_argument("--epochs", type=int) - parser.add_argument("--batch-size", type=int) - parser.add_argument("--lr", type=float) - parser.add_argument("--warmup-epochs", type=int, dest="warmup_epochs") - parser.add_argument("--hidden-dim", type=int) - parser.add_argument("--n-blocks", type=int) - parser.add_argument("--emb-dim", type=int) - parser.add_argument( - "--dropout", type=float, help="Dropout probability in ResBlocks (default: 0.1)" - ) - parser.add_argument("--val-fraction", type=float) - parser.add_argument("--seed", type=int, help="Random seed for reproducibility") - parser.add_argument( - "--validate-every", - type=int, - help="Run marginal+KL validation every N epochs (0 disables)", - ) - 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() - - cli_train = { - k: v - for k, v in { - "mode": args.mode, - "epochs": args.epochs, - "batch_size": args.batch_size, - "lr": args.lr, - "warmup_epochs": args.warmup_epochs, - "val_fraction": args.val_fraction, - "num_workers": args.num_workers, - "seed": args.seed, - "validate_every": args.validate_every, - }.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, - "dropout": args.dropout, - }.items() - if v is not None - } - 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 gconfig.auto_device() - out_dir = Path( - args.out - or ( - 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']}" - ) - ) - - print(f"device: {device}") - print(f"out_dir: {out_dir}") - - 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, - ) - - -if __name__ == "__main__": - main()