File size: 1,833 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 | #!/usr/bin/env python
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()
|