abtest_backend_gpu.py 1.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  1. #!/usr/bin/env python3
  2. # -*- coding: utf-8 -*-
  3. """gpu 后端 A/B 等价测试:同一进程先后用 milvus / gpu 后端跑同一批图,
  4. 输出(predicted_card_id + Top-1 的 id/fusion/upper/lower)必须逐字段一致。"""
  5. import os
  6. import sys
  7. sys.stdout.reconfigure(encoding="utf-8")
  8. _ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
  9. sys.path.insert(0, _ROOT)
  10. sys.path.insert(0, os.path.join(_ROOT, "scripts"))
  11. import torch
  12. from scripts.cascade_match import DualCascadeMatcher
  13. IMGS = [
  14. "data/audit_imgs/multicard_2up.jpg",
  15. "data/audit_imgs/raspi_front_2026-09-01_13-22-48_529284.jpg",
  16. "data/audit_imgs/raspi_front_2026-09-01_11-24-01_965039.jpg",
  17. ]
  18. def run(backend):
  19. dev = torch.device("cuda:0")
  20. m = DualCascadeMatcher(backend=backend, device=dev, detector_device=dev,
  21. milvus_alias="ab_" + backend)
  22. out = {}
  23. for p in IMGS:
  24. path = p if os.path.isabs(p) else os.path.join(_ROOT, p)
  25. r = m.query_all(path)
  26. out[p] = [(x["predicted_card_id"],
  27. [(t["card_id"], t["fusion"], t["upper"], t["lower"]) for t in x["top_k"]])
  28. for x in r]
  29. print(f"[{backend}] {os.path.basename(p)}: {out[p]}", flush=True)
  30. del m
  31. torch.cuda.empty_cache()
  32. return out
  33. a = run("milvus")
  34. b = run("gpu")
  35. same = a == b
  36. print(f"\n[A/B] milvus vs gpu 输出全等: {same}", flush=True)
  37. if not same:
  38. for p in a:
  39. if a[p] != b.get(p):
  40. print(f" 差异 {p}:\n milvus={a[p]}\n gpu ={b.get(p)}", flush=True)
  41. sys.exit(0 if same else 1)