File size: 2,701 Bytes
6598593 | 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 | #!/usr/bin/env python3
"""
Evaluate Venice-H1 on pre-extracted feature caches.
Usage:
python evaluate.py --checkpoint checkpoints/best.pt \
--splits data/cached_testA_refcoco_unc_feats.pt \
data/cached_testB_refcoco_unc_feats.pt
"""
import argparse
import json
from pathlib import Path
import torch
from torch.utils.data import DataLoader
from venice_h1.model import VeniceH1Reranker
from train import FeatureCacheDataset, evaluate
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--checkpoint", required=True,
help="Path to checkpoint .pt file")
parser.add_argument("--splits", nargs="+", required=True,
help="One or more feature cache .pt files to evaluate")
parser.add_argument("--tau", type=float, default=0.5,
help="Failure Gate threshold")
parser.add_argument("--no_grid", action="store_true")
parser.add_argument("--batch_size", type=int, default=512)
parser.add_argument("--device", default="cuda")
args = parser.parse_args()
device = torch.device(args.device if torch.cuda.is_available() else "cpu")
# Load checkpoint
ckpt = torch.load(args.checkpoint, map_location="cpu")
state = ckpt["model"] if "model" in ckpt else ckpt
# Infer model config from weights
use_grid = not args.no_grid
model = VeniceH1Reranker(use_grid=use_grid).to(device)
model.load_state_dict(state, strict=False)
model.eval()
print(f"Loaded: {args.checkpoint}")
print(f"Parameters: {model.num_parameters():,} | tau={args.tau}")
print()
results = {}
for split_path in args.splits:
if not Path(split_path).exists():
print(f" [skip] {split_path} not found")
continue
ds = FeatureCacheDataset(split_path, use_grid=use_grid)
loader = DataLoader(ds, batch_size=args.batch_size,
shuffle=False, num_workers=4)
split_name = Path(split_path).stem.replace("cached_", "").replace("_feats", "")
metrics = evaluate(model, loader, device, tau=args.tau)
results[split_name] = metrics
print(f" {split_name:<30s} "
f"AUC={metrics['gate_auc']:.3f} "
f"harmful={metrics['harmful_switch']*100:.3f}% "
f"correct_rerank={metrics['correct_rerank']*100:.1f}% "
f"failure_rate={metrics['failure_rate']*100:.1f}%")
print()
out_path = "eval_results.json"
with open(out_path, "w") as f:
json.dump(results, f, indent=2)
print(f"Saved results → {out_path}")
if __name__ == "__main__":
main()
|