PokemonCardSearch/code/tools/bench_fp16_gpurecall.py

#!/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()