| |
| """Evaluate the released multi-task checkpoint on one CPU or one GPU.""" |
|
|
| import argparse |
| import json |
| import os |
| import sys |
| from types import SimpleNamespace |
|
|
| SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) |
| REPO_ROOT = os.path.abspath(os.path.join(SCRIPT_DIR, "..")) |
| SRC_ROOT = os.path.join(REPO_ROOT, "src") |
| if SRC_ROOT not in sys.path: |
| sys.path.insert(0, SRC_ROOT) |
| if SCRIPT_DIR not in sys.path: |
| sys.path.insert(0, SCRIPT_DIR) |
|
|
| import torch |
|
|
| from clean_eval_utils import evaluate_clean_full, prepare_clean_records |
| from multi_task_model import MultiTaskGAT |
|
|
|
|
| def resolve_device(requested): |
| if requested == "cpu": |
| return torch.device("cpu") |
| if not torch.cuda.is_available(): |
| if requested == "cuda": |
| raise RuntimeError("CUDA was requested but is unavailable.") |
| return torch.device("cpu") |
| try: |
| probe = torch.empty(1, device="cuda:0") |
| del probe |
| return torch.device("cuda:0") |
| except RuntimeError as error: |
| if requested == "cuda": |
| raise |
| print("CUDA initialization failed; using CPU: {}".format(error), flush=True) |
| return torch.device("cpu") |
|
|
|
|
| def load_checkpoint(path, device): |
| artifact = torch.load(path, map_location="cpu") |
| if not isinstance(artifact, dict) or artifact.get("artifact_format") != "icflownet_model_v1": |
| raise RuntimeError("Expected an icflownet_model_v1 checkpoint.") |
| config = dict(artifact.get("model_config") or {}) |
| config.pop("architecture_version", None) |
| model = MultiTaskGAT(**config) |
| model.load_state_dict(artifact["model_state_dict"], strict=True) |
| model.to(device) |
| model.eval() |
| return model, config |
|
|
|
|
| def evaluate_scope(args, scope, device): |
| clean_root = os.path.abspath(args.clean_root) |
| records, preflight = prepare_clean_records( |
| clean_root, os.path.join(clean_root, "resources"), scope |
| ) |
| if args.max_records is not None: |
| records = records[: args.max_records] |
| model, model_config = load_checkpoint(os.path.abspath(args.checkpoint), device) |
| output_dir = os.path.join(os.path.abspath(args.output_dir), scope) |
| os.makedirs(output_dir, exist_ok=True) |
| metrics = evaluate_clean_full( |
| SimpleNamespace(module=model), |
| records, |
| device, |
| rank=0, |
| world_size=1, |
| output_dir=output_dir, |
| log_every=args.log_every, |
| amp_enabled=False, |
| ) |
| result = { |
| "format": "icflownet_cleantest_eval_v1", |
| "scope": scope, |
| "checkpoint": os.path.abspath(args.checkpoint), |
| "checkpoint_artifact_format": "icflownet_model_v1", |
| "threshold": 0.5, |
| "clean_root": clean_root, |
| "world_size": 1, |
| "device": str(device), |
| "precision": "fp32", |
| "amp_enabled": False, |
| "model_config": model_config, |
| "preflight": preflight, |
| "metrics": metrics, |
| } |
| result_path = os.path.join(output_dir, "result.json") |
| with open(result_path, "w") as handle: |
| json.dump(result, handle, indent=2) |
| handle.write("\n") |
| print(json.dumps({"result": result_path, "scope": scope, "metrics": metrics}, indent=2)) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--checkpoint", required=True) |
| parser.add_argument("--clean_root", required=True) |
| parser.add_argument("--scope", choices=("overall", "long_range", "both"), default="both") |
| parser.add_argument("--output_dir", required=True) |
| parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto") |
| parser.add_argument("--max_records", type=int) |
| parser.add_argument("--log_every", type=int, default=10) |
| args = parser.parse_args() |
| if args.max_records is not None and args.max_records < 1: |
| parser.error("--max_records must be at least 1") |
| device = resolve_device(args.device) |
| scopes = ("overall", "long_range") if args.scope == "both" else (args.scope,) |
| for scope in scopes: |
| evaluate_scope(args, scope, device) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
|
|
