File size: 2,382 Bytes
31e2456 | 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 | """CLI entry point: train a single model variant."""
from __future__ import annotations
import argparse
import os
from dotenv import load_dotenv
load_dotenv()
os.environ.setdefault("HF_TOKEN", os.environ.get("HUGGINGFACE_API_KEY", ""))
os.environ.setdefault("WANDB_API_KEY", os.environ.get("WANDB_API_KEY", ""))
from physiojepa.trainer import TrainConfig, load_yaml_config, train
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--config", required=True)
ap.add_argument("--run_name", type=str, default=None)
ap.add_argument("--model", type=str, default=None, choices=["A", "B", "C", "F"])
ap.add_argument("--epochs", type=int, default=None)
ap.add_argument("--batch_size", type=int, default=None)
ap.add_argument("--index_path", type=str, default=None)
ap.add_argument("--shard_roots_json", type=str, default=None,
help="JSON file listing shard roots")
ap.add_argument("--wandb_mode", type=str, default=None)
ap.add_argument("--num_workers", type=int, default=None)
ap.add_argument("--output_dir", type=str, default=None)
ap.add_argument("--subset_frac", type=float, default=None)
ap.add_argument("--log_every", type=int, default=None)
ap.add_argument("--ema_start", type=float, default=None)
ap.add_argument("--ema_end", type=float, default=None)
ap.add_argument("--ema_warmup_frac", type=float, default=None)
ap.add_argument("--seed", type=int, default=None)
ap.add_argument("--pred_depth", type=int, default=None)
ap.add_argument("--query_mode", type=str, default=None, choices=["learned", "sinusoidal"])
ap.add_argument("--mask_ratio", type=float, default=None)
ap.add_argument("--fast_cache_dir", type=str, default=None)
args = ap.parse_args()
cfg = load_yaml_config(args.config)
overrides = {k: v for k, v in vars(args).items() if v is not None and k not in ("config",)}
if "shard_roots_json" in overrides:
import json
cfg.shard_roots = json.loads(open(overrides.pop("shard_roots_json")).read())
for k, v in overrides.items():
setattr(cfg, k, v)
print(f"[train] resolved config: model={cfg.model} run={cfg.run_name} "
f"epochs={cfg.epochs} bs={cfg.batch_size} shards={len(cfg.shard_roots)}")
res = train(cfg)
print(f"[train] done: {res}")
if __name__ == "__main__":
main()
|