train.py

"""Train one configured run.

    python -O -m src.train --config configs/segthor_enet_ce.yaml
    python -O -m src.train --config configs/segthor_enet_ce.yaml --set seed=1 optim.kwargs.lr=0.001
    python -m src.train --config configs/segthor_enet_ce.yaml --smoke     # 2 epochs, 16 slices
    python -O -m src.train --config configs/full_cv4_enet_ce.yaml --fold 2   # one fold of a CV config

Re-running the same command resumes an interrupted run, or exits if it already finished.
"""
import argparse
import csv
import traceback
from pathlib import Path

import torch

from src.checkpoint import seed_everything
from src.config import config_hash, load_config
from src.data import ensure_sliced, has_val
from src.engine import fit
from src.folds import run_configs
from src.plots import plot_curves
from src.run import (copy_back, environment, now, prepare_run_dir, read_json, resolve_device,
                     setup_logging, update_manifest, write_json)
from src.wandb_logger import WandbLogger


def summarize(run_dir: Path, cfg: dict, best_epoch: int) -> dict:
    with (run_dir / "epochs.csv").open(newline="") as f:
        rows = list(csv.DictReader(f))
    best = next(r for r in rows if int(r["epoch"]) == best_epoch)
    manifest = read_json(run_dir / "manifest.json")
    return {"experiment": cfg["experiment"], "seed": cfg["seed"], "run": run_dir.name,
            "config_hash": manifest["config_hash"], "model": cfg["model"]["name"],
            "loss": cfg["loss"]["name"], "optim": cfg["optim"]["name"],
            "lr": cfg["optim"]["kwargs"].get("lr"), "data_root": cfg["data"]["root"],
            "epochs": len(rows), "best_epoch": best_epoch,
            "best": {k: float(v) for k, v in best.items() if k.startswith(("val_", "train_"))},
            "user": manifest["user"], "git_commit": manifest["git_commit"],
            "started": manifest["started"], "finished": now(), "wandb_url": manifest.get("wandb_url")}


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE", help="override config values")
    parser.add_argument("--smoke", action="store_true", help="tiny run to catch errors before using GPU time")
    parser.add_argument("--force", action="store_true", help="move an existing run aside and start over")
    parser.add_argument("--device", choices=("cuda", "cpu"), help="override the config's `device`")
    parser.add_argument("--fold", type=int, help="run only this zero-based fold of a cross-validation config")
    args = parser.parse_args(argv)

    overrides = list(args.set) + ([f"device={args.device}"] if args.device else [])
    cfg = load_config(args.config, overrides, smoke=args.smoke)
    # a cross-validation config (data.preprocess.fold > 1) runs its folds one after another, or only --fold
    for run_cfg in run_configs(cfg, args.fold, args.smoke):
        train_one(run_cfg, force=args.force, smoke=args.smoke)


def train_one(cfg: dict, force: bool = False, smoke: bool = False) -> None:
    """Train or resume one split with its own model, optimizer and outputs."""
    prepared = prepare_run_dir(cfg, force=force, smoke=smoke)
    if prepared is None:
        print(f"{cfg['experiment']} seed {cfg['seed']} already finished; use --force to re-run")
        return
    run_dir, resume = prepared
    log = setup_logging(run_dir / "train.log")
    device = resolve_device(cfg["device"])

    manifest = read_json(run_dir / "manifest.json")
    if resume:
        manifest = update_manifest(run_dir, status="running",
                                   resumes=manifest.get("resumes", []) + [environment(device) | {"at": now()}])
    else:
        manifest = update_manifest(run_dir, status="running", started=now(), config_hash=config_hash(cfg),
                                   smoke=smoke, **environment(device))
    log.info("%s %s -> %s on %s", "resuming" if resume else "starting", cfg["experiment"], run_dir, device)

    wb = WandbLogger(cfg, run_dir, run_name=f"{cfg['experiment']}/seed{cfg['seed']}",
                     run_id=manifest.get("wandb_id"))
    if wb.id:
        update_manifest(run_dir, wandb_id=wb.id, wandb_url=wb.url)
    seed_everything(cfg["seed"])  # after W&B init, which uses `random` itself; resume restores RNG later
    try:
        ensure_sliced(cfg)
        best_epoch = fit(cfg, run_dir, device, resume, wb)
    except BaseException as exc:
        log.error("run failed:\n%s", traceback.format_exc())
        update_manifest(run_dir, status="failed", failed_at=now(), error=repr(exc))
        wb.finish(exit_code=1)
        raise

    summary = summarize(run_dir, cfg, best_epoch)
    write_json(run_dir / "summary.json", summary)
    prefix = "val" if has_val(cfg) else "train"  # a final run (retains 0) has no val split
    dice_cols = [f"{prefix}_dice_{cfg['data']['class_names'][k]}" for k in cfg["eval"]["classes"]] + [f"{prefix}_dice_fg"]
    plot_curves(run_dir / "epochs.csv", run_dir / "plots" / "curves.png", dice_cols)
    update_manifest(run_dir, status="completed", finished=summary["finished"])
    wb.summary({f"best/{k}": v for k, v in summary["best"].items()} | {"best_epoch": best_epoch})
    wb.finish()
    dest = copy_back(run_dir, cfg)
    (run_dir / ".done").touch()
    log.info("done: best epoch %d, %s dice fg %.4f; results in %s%s", best_epoch, prefix,
             summary["best"][f"{prefix}_dice_fg"], run_dir, f", copied to {dest}" if dest else "")


if __name__ == "__main__":
    main()