#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 8020 优化前验证：DINOv2 FP16 autocast + GPU 常驻图库召回（249 上跑，只测不改生产）。 Part A DINOv2 FP16 autocast（monkey-patch amp 版 _forward） A1 单卡(bs=1) 前向耗时：FP32 vs AMP（各 N 次，含 .cpu().numpy() 完整路径） A2 特征一致性：逐 half 余弦(fp32, fp16) A3 端到端等价：两种特征各自走完 milvus 检索+rank，比对 Top-5 (card_id, fusion) A4 多卡批收益：bs=2 一次前向 vs 两次 bs=1 Part B GPU 常驻图库召回（npy 图库 → 显存 torch 矩阵乘，替代 milvus 网络往返） B1 一致性：milvus search vs GPU fp32 vs GPU fp16 的 Top-30 重叠 + sim 最大差 B2 融合一致性：三种 cand 喂 rank(α) 后 Top-5 是否一致 B3 耗时：单卡检索(milvus vs GPU32 vs GPU16) ms B4 生产组合投影：query_all 端到端（原版 / 仅AMP / AMP+GPU召回）各 N 次 用法（pytorch env，249，与 8020 共存只占显存不冲突——8020 空闲时测）: CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=1 \ ~/miniconda3/envs/pytorch/bin/python tools/bench_fp16_gpurecall.py --n 50 """ import os import sys import time import argparse import statistics sys.stdout.reconfigure(encoding="utf-8") _ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, _ROOT) sys.path.insert(0, os.path.join(_ROOT, "scripts")) import numpy as np from PIL import Image import config from scripts.cascade_match import DualCascadeMatcher from modules.feature_extractor_dual import letterbox392, split_half, half_to_chw def _amp_forward(self, backbone, pil_halves): """生产将采用的 AMP 版 _forward：autocast fp16 前向，cls.float() 升回 fp32 再 L2。 与原版逻辑逐行对应，仅 forward 段不同。""" import torch import torch.nn.functional as F out = np.full((len(pil_halves), self.feature_dim), np.nan, dtype=np.float32) tensors, valid = [], [] for i, p in enumerate(pil_halves): if p is None: continue try: tensors.append(torch.from_numpy(half_to_chw(p)).float()) valid.append(i) except Exception: continue if not tensors: return out for s in range(0, len(tensors), self.batch_size): chunk = torch.stack(tensors[s:s + self.batch_size]).to(self.device) with torch.no_grad(): try: with torch.autocast("cuda", dtype=torch.float16): try: cls = backbone(chunk, interpolate_pos_encoding=True).last_hidden_state[:, 0, :] except TypeError: cls = backbone(chunk).last_hidden_state[:, 0, :] cls = cls.float() # fp16 累加结果升回 fp32 再归一化 except RuntimeError: raise feat = F.normalize(cls, p=2, dim=1).cpu().numpy() for k, vi in enumerate(valid[s:s + self.batch_size]): out[vi] = feat[k] return out def _timed(fn, n): """跑 n 次返回 mean/p50 ms。""" for _ in range(3): fn() ts = [] for _ in range(n): t0 = time.perf_counter() fn() ts.append((time.perf_counter() - t0) * 1000) return statistics.mean(ts), statistics.median(ts) def _top5(matcher, q_u, q_l, backend_gpu=None): """q 特征 → 检索 + rank → Top-5 [(card_id, round(fusion,4))...]。 backend_gpu: None=matcher 自带 backend（milvus）；否则 ('gpu', G_u, G_l, ids)""" if backend_gpu is None: cand_ids, su, sl = matcher._query_raw_milvus(q_u, q_l) else: _, Gu, Gl, ids = backend_gpu import torch q = torch.from_numpy(q_u).to(Gu.device).to(Gu.dtype) # 图库 fp16 时查询向量同步转 fp16 sims = Gu @ q K = min(matcher.top_k_recall, Gu.shape[0]) tv = torch.topk(sims, K) cand = tv.indices.cpu().numpy() cand_ids = [ids[int(c)] for c in cand] su = tv.values.float().cpu().numpy() ql = torch.from_numpy(q_l).to(Gl.device).to(Gl.dtype) sl = (Gl[tv.indices] @ ql).float().cpu().numpy() raw = {"cand_ids": cand_ids, "upper_sim": su, "lower_sim": sl} ranked = matcher.rank(raw, matcher.alpha) return [(r[0], round(r[1], 4)) for r in ranked[:5]] def main(): ap = argparse.ArgumentParser() ap.add_argument("--images", nargs="+", default=[ "data/audit_imgs/raspi_front_2026-09-01_13-22-48_529284.jpg", "data/audit_imgs/multicard_2up.jpg"]) ap.add_argument("--n", type=int, default=50) ap.add_argument("--device", default="cuda:0") args = ap.parse_args() import torch dev = torch.device(args.device) print(f"[Bench] 初始化 matcher（device={dev}）...", flush=True) matcher = DualCascadeMatcher( alpha=config.CASCADE_ALPHA, top_k_recall=config.CASCADE_TOP_K_RECALL, device=dev, detector_device=dev, milvus_alias="prof2", verbose=True) ex = matcher.extractor # ---- 准备真实输入：crops → halves ---- crops = [] for p in args.images: path = p if os.path.isabs(p) else os.path.join(_ROOT, p) arr = np.array(Image.open(path).convert("RGB")) for inst in matcher.detector.detect_and_crop_all(arr): crops.append(inst["crop"]) print(f"[Bench] 测试输入：{len(crops)} 张真实裁卡", flush=True) halves = [] half_keys = [] # 与 halves 平行，存入 dict 时用的 key（PIL 重建后 id 会变，必须先存） for c in crops: lb = letterbox392(Image.fromarray(c)) up, lo = split_half(lb) halves.append(("up", up)) halves.append(("lo", lo)) half_keys.append(("up", id(up))) half_keys.append(("lo", id(lo))) # ==================== Part A：DINOv2 FP16 ==================== print(f"\n========== Part A DINOv2 FP16 autocast ==========", flush=True) # A1/A2 逐 half 特征与耗时 fp32_feats = {} for kind, img in halves: fn = ex.extract_upper if kind == "up" else ex.extract_lower fp32_feats[(kind, id(img))] = fn([img])[0] def run_fp32_all(): for kind, img in halves: fn = ex.extract_upper if kind == "up" else ex.extract_lower fn([img]) m32, p32 = _timed(run_fp32_all, args.n) per_half_32 = m32 / len(halves) print(f"A1 FP32 单卡请求(2 前向 bs=1): mean={m32:.1f}ms p50={p32:.1f}ms " f"→ 每 half {per_half_32:.1f}ms", flush=True) # 切 AMP orig_forward = ex._forward ex._forward = _amp_forward.__get__(ex) amp_feats = {} for kind, img in halves: fn = ex.extract_upper if kind == "up" else ex.extract_lower amp_feats[(kind, id(img))] = fn([img])[0] def run_amp_all(): for kind, img in halves: fn = ex.extract_upper if kind == "up" else ex.extract_lower fn([img]) m16, p16 = _timed(run_amp_all, args.n) per_half_16 = m16 / len(halves) print(f"A1 AMP 单卡请求(2 前向 bs=1): mean={m16:.1f}ms p50={p16:.1f}ms " f"→ 每 half {per_half_16:.1f}ms 加速比 {m32/m16:.2f}x", flush=True) cos_list = [] for key in fp32_feats: a, b = fp32_feats[key], amp_feats[key] if np.isnan(a).any() or np.isnan(b).any(): cos_list.append(float("nan")) continue cos_list.append(float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))) cos_list = [c for c in cos_list if not np.isnan(c)] print(f"A2 特征余弦(fp32 vs fp16): min={min(cos_list):.6f} " f"mean={statistics.mean(cos_list):.6f} n={len(cos_list)}/{len(halves)}", flush=True) # A3 端到端等价：两种特征走完整 milvus 检索+rank diff_cnt = 0 for ci in range(len(crops)): t5a = _top5(matcher, fp32_feats[half_keys[2 * ci]], fp32_feats[half_keys[2 * ci + 1]]) t5b = _top5(matcher, amp_feats[half_keys[2 * ci]], amp_feats[half_keys[2 * ci + 1]]) if t5a != t5b: diff_cnt += 1 print(f" [A3差异] crop{ci}: fp32={t5a[:2]} vs amp={t5b[:2]}", flush=True) print(f"A3 Top-5 端到端一致性(milvus): {len(crops)-diff_cnt}/{len(crops)} 完全一致", flush=True) # A4 多卡批收益（用双卡图的 2 个 upper halves 合一批） ups = [img for kind, img in halves if kind == "up"][:2] if len(ups) == 2: m_b1, _ = _timed(lambda: (ex.extract_upper([ups[0]]), ex.extract_upper([ups[1]])), args.n) m_b2, _ = _timed(lambda: ex.extract_upper(ups), args.n) print(f"A4 批量: 2次bs=1={m_b1:.1f}ms vs 1次bs=2={m_b2:.1f}ms " f"(每 half {m_b1/2:.1f} → {m_b2/2:.1f}ms)", flush=True) # 还原 FP32（Part B 用两种特征分别测，最后再切回） ex._forward = orig_forward # ==================== Part B：GPU 常驻图库 ==================== print(f"\n========== Part B GPU 常驻图库召回 ==========", flush=True) t0 = time.perf_counter() Gu32 = torch.from_numpy(np.load(config.GALLERY_UPPER_FEATURES_PATH)).to(dev) Gl32 = torch.from_numpy(np.load(config.GALLERY_LOWER_FEATURES_PATH)).to(dev) ids = matcher.card_ids print(f"B0 图库上卡: {tuple(Gu32.shape)} {time.perf_counter()-t0:.1f}s " f"显存 +{(Gu32.element_size()*Gu32.nelement()+Gl32.element_size()*Gl32.nelement())/2**20:.0f}MB", flush=True) assert Gu32.shape[0] == len(ids) gpu32 = ("gpu", Gu32, Gl32, ids) Gu16, Gl16 = Gu32.half(), Gl32.half() gpu16 = ("gpu", Gu16, Gl16, ids) # 一致性：对每个 crop 的 fp32 特征，三路各出 Top-5 def bench_backends(q_u, q_l, tag): t5_mil = _top5(matcher, q_u, q_l) t5_g32 = _top5(matcher, q_u, q_l, gpu32) t5_g16 = _top5(matcher, q_u, q_l, gpu16) ok32 = t5_mil == t5_g32 ok16 = t5_mil == t5_g16 print(f" [{tag}] milvusTop1={t5_mil[0]} gpu32一致={ok32} gpu16一致={ok16}", flush=True) if not ok32: print(f" gpu32={t5_g32[:2]}", flush=True) if not ok16: print(f" gpu16={t5_g16[:2]}", flush=True) return ok32, ok16 n_ok32 = n_ok16 = 0 for c in crops: lb = letterbox392(Image.fromarray(c)) up, lo = split_half(lb) q_u = ex.extract_upper([up])[0] q_l = ex.extract_lower([lo])[0] a, b = bench_backends(q_u, q_l, "fp32特征") n_ok32 += int(a) n_ok16 += int(b) print(f"B1/B2 Top-5 一致(milvus 基准): gpu32 {n_ok32}/{len(crops)} gpu16 {n_ok16}/{len(crops)}", flush=True) # B3 耗时（纯检索段，单卡一次） q_u0, q_l0 = None, None lb = letterbox392(Image.fromarray(crops[0])) up0, lo0 = split_half(lb) q_u0 = ex.extract_upper([up0])[0] q_l0 = ex.extract_lower([lo0])[0] m_mil, _ = _timed(lambda: matcher._query_raw_milvus(q_u0, q_l0), 100) def gpu_call(Gu, Gl, q_u, q_l): sims = Gu @ q_u tv = torch.topk(sims, matcher.top_k_recall) cand = tv.indices sl = Gl[cand] @ q_l return tv.values, sl q_uT = torch.from_numpy(q_u0).to(dev) q_lT = torch.from_numpy(q_l0).to(dev) m_g32, _ = _timed(lambda: gpu_call(Gu32, Gl32, q_uT, q_lT), 200) m_g16, _ = _timed(lambda: gpu_call(Gu16, Gl16, q_uT.half(), q_lT.half()), 200) print(f"B3 单卡检索段: milvus={m_mil:.1f}ms gpu32={m_g32:.2f}ms gpu16={m_g16:.2f}ms", flush=True) # B4 生产组合投影：query_all 端到端 print(f"\n---------- B4 query_all 端到端投影（双卡图） ----------", flush=True) img2 = os.path.join(_ROOT, "data/audit_imgs/multicard_2up.jpg") def patch_gpu_recall(m): """临时替换 query_all 内的检索分支 → GPU 常驻图库（fp16）。""" Gu, Gl, cids = Gu16, Gl16, ids def _query_raw_gpu(self, q_u, q_l): qt = torch.from_numpy(q_u).to(dev).half() ql = torch.from_numpy(q_l).to(dev).half() sims = Gu @ qt K = min(self.top_k_recall, Gu.shape[0]) tv = torch.topk(sims, K) cand = tv.indices.cpu().numpy() cand_ids = [cids[int(c)] for c in cand] su = tv.values.float().cpu().numpy() sl = (Gl[torch.from_numpy(cand).to(dev)] @ ql).float().cpu().numpy() return cand_ids, su, sl import scripts.cascade_match as cm cm.DualCascadeMatcher._query_raw_gpu = _query_raw_gpu orig = cm.DualCascadeMatcher._query_raw_milvus def _dispatch(self, q_u, q_l): return self._query_raw_gpu(q_u, q_l) m._query_raw_milvus = _dispatch.__get__(m) return orig # 基线 mb0, pb0 = _timed(lambda: matcher.query_all(img2), args.n) print(f"B4 原版(milvus+FP32) : mean={mb0:.0f}ms", flush=True) # 仅 AMP ex._forward = _amp_forward.__get__(ex) mb1, pb1 = _timed(lambda: matcher.query_all(img2), args.n) print(f"B4 仅 AMP : mean={mb1:.0f}ms ({mb0/mb1:.2f}x)", flush=True) # AMP + GPU 召回 orig_milvus = patch_gpu_recall(matcher) mb2, pb2 = _timed(lambda: matcher.query_all(img2), args.n) print(f"B4 AMP + GPU召回(fp16图库) : mean={mb2:.0f}ms ({mb0/mb2:.2f}x)", flush=True) # 等价性：最终输出比对 r0 = matcher.query_all(img2) out0 = [(r["predicted_card_id"], r["top_k"]) for r in r0] ex._forward = orig_forward m._query_raw_milvus = orig_milvus.__get__(matcher) r1 = matcher.query_all(img2) out1 = [(r["predicted_card_id"], r["top_k"]) for r in r1] same = out0 == out1 print(f"B4 等价性(优化后 vs 原版 输出全等): {same}", flush=True) if not same: for a, b in zip(out0, out1): if a != b: print(f" 原版={a}\n 优化={b}", flush=True) print(f"\n[Bench 完成] 预期单卡请求: 176ms → 约{176 - (per_half_32-per_half_16)*2 - (m_mil-m_g16):.0f}ms" f"（DINO省{(per_half_32-per_half_16)*2:.0f} + 检索省{m_mil-m_g16:.0f}）", flush=True) if __name__ == "__main__": main()