| |
| from __future__ import annotations |
| import argparse |
| import json |
| from pathlib import Path |
| import torch |
| from sacflow.utils.config import load_yaml |
| from sacflow.utils.misc import seed_everything, ensure_dir |
| from sacflow.utils.distributed import init_distributed, cleanup, is_main_process |
| from sacflow.data.loader import build_loader |
| from sacflow.models.unet3d import build_model |
| from sacflow.engine.train_loop import evaluate |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--config", required=True) |
| ap.add_argument("--checkpoint", default=None) |
| ap.add_argument("--split", default=None) |
| args = ap.parse_args() |
| cfg = load_yaml(args.config) |
| if args.checkpoint: |
| cfg.setdefault("eval", {})["checkpoint"] = args.checkpoint |
| if args.split: |
| cfg.setdefault("eval", {})["split"] = args.split |
| seed_everything(int(cfg.get("seed", 1337))) |
| device = init_distributed(cfg.get("distributed", {}).get("backend", "nccl")) |
| model = build_model(cfg).to(device) |
| ckpt_path = cfg.get("eval", {}).get("checkpoint") or cfg.get("train", {}).get("source_checkpoint") |
| if ckpt_path is None: |
| ckpt_path = str(Path(cfg["output_dir"]) / "checkpoints" / "best.pt") |
| ckpt = torch.load(ckpt_path, map_location="cpu") |
| model.load_state_dict(ckpt.get("model", ckpt), strict=False) |
| split = cfg.get("eval", {}).get("split", "target_test") |
| loader = build_loader(cfg, split=split, training=False, require_label=True) |
| metrics = evaluate(model, loader, cfg, device) |
| if is_main_process(): |
| print(json.dumps(metrics, indent=2)) |
| out = ensure_dir(Path(cfg["output_dir"]) / "eval") |
| with open(out / f"metrics_{split}.json", "w") as f: |
| json.dump(metrics, f, indent=2) |
| cleanup() |
|
|
| if __name__ == "__main__": |
| main() |
|
|
