File size: 4,107 Bytes
6de06a8 | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | #!/usr/bin/env python3
"""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()
|