ann_cost.py

1"""2ANN 的成本直觉：为什么要做近似搜索（第18章 §18.1）34正文给的算例：物品库 N=10^9、向量维度 d=128，一次精确检索约需5    N × d ≈ 10^11 次浮点运算6即使在高性能服务器上也难以满足毫秒级响应。78本脚本做四件事：91) 复现这个算例，并把它换算成"单核每秒 10 亿次浮点运算"下的耗时与吞吐；102) 在本机真实测量 NumPy 暴力检索的吞吐，用实测值再算一遍（数字更具体）；113) 计算存储成本：float32 / int8 / PQ 三种编码下 10^9 个向量各占多少内存；124) 打印"只扫描 x% 的库、召回率还能保持多少"的权衡表——这是所有 ANN 算法13   共同的 trade-off，后面的 KD 树、PQ、HNSW 只是用不同方式实现这个扫描比例。1415[MOCK]：机器算力（10 GFLOPs）、内存带宽、内存价格按典型云主机量级取整。16"""1718from __future__ import annotations1920import pathlib21import sys2223import numpy as np2425sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent))2627from common import (  # noqa: E40228    knn_exact, make_queries, make_vectors, recall_at_k, sq_dists, table, timing,29    topk_from_sq,30)3132N_DB = 10 ** 9          # 物品库规模33DIM = 128               # Embedding 维度34TOPK = 10035GFLOP_PER_SEC = 10.0    # [MOCK] 单核算力，GFLOP/s3637# 本机实测用的小库38N_MEASURED = 20_00039N_QUERIES = 200404142def main() -> None:43    print("ANN 的成本直觉（§18.1）\n")4445    # ---------------- 1) 正文算例 ----------------46    flops = N_DB * DIM47    print("① 精确检索的算力账（正文 N=10^9, d=128）：")48    print(f"  单次查询距离计算量 : N × d = {N_DB:,} × {DIM} = {flops:.1e} FLOPs")49    print(f"  平方距离约为 3 次乘加 ≈ {3 * flops:.1e} FLOPs（|x|²+|y|²-2xᵀy 展开）")50    print(f"  按 {GFLOP_PER_SEC:.0f} GFLOP/s 单核估算单次查询耗时 : "51          f"{flops / (GFLOP_PER_SEC * 1e9):.1f} s")52    print(f"  一个 32 核的检索节点、单次查询仍要 "53          f"{flops / (GFLOP_PER_SEC * 1e9 * 32) * 1000:.0f} ms")54    print("  → 毫秒级 SLA 下只能接受访问 10^4–10^5 个候选，即全库的 0.001%–0.01%。\n")5556    # ---------------- 2) 本机实测 ----------------57    X, _ = make_vectors(n=N_MEASURED, dim=DIM, n_clusters=50, spread=0.4)58    Q = make_queries(X, n_queries=N_QUERIES)59    t_brute, d = timing(sq_dists, X, Q, repeat=3)60    qps = N_QUERIES / t_brute61    print("② 本机 NumPy 暴力检索实测（BLAS 单机多线程）：")62    print(f"  库规模 {N_MEASURED:,} × {DIM} 维，{N_QUERIES} 个查询，耗时 {t_brute * 1000:.1f} ms")63    print(f"  吞吐 ≈ {qps:,.0f} QPS，单查询 ≈ {t_brute / N_QUERIES * 1e6:.0f} µs")64    print(f"  按这个吞吐外推到 10^9 条：单查询 ≈ "65          f"{t_brute / N_QUERIES * (N_DB / N_MEASURED) * 1000:.0f} ms")66    print("  注意：实测值包含了 BLAS 的多线程与缓存命中，不代表「真正的单核算力」；")67    print("  外推值只能说明数量级——即便乐观估计，也与毫秒级 SLA 相差 3 个数量级以上。\n")6869    # ---------------- 3) 存储 ----------------70    print("③ 存储账（10^9 条向量）：")71    rows = []72    encodings = [73        ("float32", 4 * DIM, "原始精度，Faiss IndexFlat / HNSWFlat"),74        ("float16", 2 * DIM, "半精度，召回略降，HNSW 常用"),75        ("int8 (SQ)", DIM, "标量量化，每维 8 bit"),76        ("PQ m=8,K=256", 8, "8 个子空间 × 8 bit，§18.4 详述"),77        ("PQ m=16,K=256", 16, "16 个子空间 × 8 bit，精度更高"),78        ("二值码 128bit", 16, "1 bit/维，汉明距离检索"),79    ]80    for name, nbytes, note in encodings:81        total_gb = N_DB * nbytes / 1e982        rows.append([name, f"{nbytes} B/条", f"{total_gb:,.1f} GB", note])83    table(["编码", "单条大小", "总量", "典型用途"], rows)84    print("\n  再算一笔账：10^9 条 float32 正好 512 GB，一台 512 GB 内存的机器放完向量就")85    print("  没有余量给索引结构、操作系统与副本了。所以十亿级要么压缩（PQ/二值码），")86    print("  要么分片（§18.7 的分布式 ANN），二者通常一起用。\n")8788    # ---------------- 4) 扫描比例 vs 召回 ----------------89    # 用"随机采样候选 + 精确重排"模拟：只访问 p 比例的库，能保住多少 Recall@10090    truth, _ = knn_exact(X, Q, TOPK)91    rng = np.random.default_rng(0)92    print("④ 扫描比例 vs 召回（随机采样 p 比例的库向量，再精确重排）：")93    rows = []94    for p in (1.0, 0.5, 0.1, 0.01, 0.001):95        m = max(TOPK, int(N_MEASURED * p))96        idx = np.stack([rng.choice(N_MEASURED, m, replace=False) for _ in Q])97        sub = np.asarray([[d[i, j] for j in row] for i, row in enumerate(idx)])98        pred, _ = topk_from_sq(sub, TOPK)99        pred = np.take_along_axis(idx, pred, axis=1)100        r = recall_at_k(pred, truth)101        rows.append([f"{p:.1%}", f"{m:,}", f"{r:.3f}",102                     "暴力" if p == 1.0 else "无结构采样，只是下界"])103    table(["扫描比例", "访问候选数", "Recall@100", "说明"], rows)104    print("\n  这张表说明：随机丢候选的 Recall 是线性掉的，1% 的扫描只剩约 1% 召回。")105    print("  ANN 的全部价值在于——用结构把「该扫的那 1%」找出来，而不是均匀地扫 1%。")106    print("  后面的 KD 树(§18.2)、LSH(§18.3)、PQ(§18.4)、HNSW(§18.5) 都是在做这件事。\n")107108    # ---------------- 5) 复杂度对照 ----------------109    print("⑤ 各算法一次查询的复杂度（N=库大小，d=维度）：")110    table(111        ["算法", "查询复杂度", "依赖的结构", "主要误差来源"],112        [113            ["精确检索", "O(N·d)", "无", "无"],114            ["KD 树 / 球树", "低维 O(log N)，高维退化 O(N)", "超矩形 / 超球", "剪枝失效（维度灾难）"],115            ["随机投影树", "O(log N)（内在维度低时）", "随机超平面", "随机划分的方差"],116            ["LSH", "O(N·(1-ε)) 亚线性", "哈希桶", "碰撞概率 / 桶内还需精确算"],117            ["PQ + IVF", "O(K·d + m·N/nprobe)", "码本 + 倒排表", "量化误差"],118            ["HNSW", "O(efSearch · log N)", "多层近邻图", "图不连通 / 局部最优"],119        ],120    )121122123if __name__ == "__main__":124    main()