File size: 1,857 Bytes
93ffd19 | 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 | #!/usr/bin/env python
from __future__ import annotations
import argparse
from pathlib import Path
import yaml
from sacflow.utils.config import load_yaml, save_yaml, set_by_path
from sacflow.utils.misc import seed_everything, ensure_dir
from sacflow.utils.distributed import init_distributed, cleanup, is_main_process, barrier
from sacflow.utils.wandb_utils import init_wandb, wandb_finish
from sacflow.engine.train_loop import run_training
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--config", required=True)
ap.add_argument("--name", default=None)
ap.add_argument("--resume", nargs="?", const="auto", default=None, help="Resume checkpoint path, or use --resume without a value for auto-detection.")
ap.add_argument("--opts", nargs="*", default=[], help="Override config values: key=value, e.g. train.epochs=220 optim.lr=1e-4")
args = ap.parse_args()
cfg = load_yaml(args.config)
for opt in args.opts:
if "=" not in opt:
raise ValueError(f"Invalid --opts entry {opt!r}; expected key=value")
key, value = opt.split("=", 1)
try:
parsed = yaml.safe_load(value)
except Exception:
parsed = value
set_by_path(cfg, key, parsed)
if args.resume is not None:
cfg.setdefault("train", {})["resume_checkpoint"] = args.resume
seed_everything(int(cfg.get("seed", 1337)))
device = init_distributed(cfg.get("distributed", {}).get("backend", "nccl"))
out_dir = ensure_dir(cfg["output_dir"])
if is_main_process():
save_yaml(cfg, out_dir / "resolved_config.yaml")
run = init_wandb(cfg, run_name=args.name or Path(cfg["output_dir"]).name)
try:
run_training(cfg, device, wandb_run=run)
finally:
barrier()
wandb_finish(run)
cleanup()
if __name__ == "__main__":
main()
|