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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
| import math
import os
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
import numpy as np
import psycopg2
from pgvector.psycopg2 import register_vector
from torchvision import datasets
PG_URL = "postgresql://postgres:yourpassword@localhost:5432/testdb"
# 每个线程一份独立的数据库连接 (psycopg2 连接/游标线程不安全)
_local = threading.local()
_conn_lock = threading.Lock()
_all_conns = []
def get_conn():
conn = getattr(_local, "conn", None)
if conn is None or conn.closed:
conn = psycopg2.connect(PG_URL)
# 让 psycopg2 能识别 vector 类型
register_vector(conn)
_local.conn = conn
with _conn_lock:
_all_conns.append(conn)
return conn
# 使用手写数字的测试集
mnist_test = datasets.MNIST(root="./data", train=False, download=True)
def l2_normalize(vec):
norm = math.sqrt(sum(v * v for v in vec))
return vec if norm == 0.0 else [v / norm for v in vec]
def topk_by_column(cur, query_vec, column, k=5):
"""在指定列上做欧氏距离最近邻检索,返回 [(label, distance), ...]。"""
# <-> 欧氏距离运算符
cur.execute(
f"""
SELECT label, {column} <-> %s::vector AS distance
FROM mnist_images
ORDER BY distance
LIMIT %s
""",
(query_vec, k),
)
return cur.fetchall()
def process_one(idx, k):
"""处理单条样本,返回 6 个统计量: (raw_top1, raw_top5, raw_dist,
emb_top1, emb_top5, emb_dist)。每个线程使用各自的连接与游标。"""
cur = get_conn().cursor()
try:
image, true_label = mnist_test[idx]
# 0~255 原始像素
arr = np.array(image).flatten().tolist()
# 归一化向量
emb = l2_normalize(arr)
# 1) 基于原始向量 pixels
raw_res = topk_by_column(cur, arr, "pixels", k)
raw_top1 = int(raw_res[0][0] == true_label)
raw_top5 = int(any(lbl == true_label for lbl, _ in raw_res))
raw_dist = raw_res[0][1]
# 2) 基于归一化向量 embedding
emb_res = topk_by_column(cur, emb, "embedding", k)
emb_top1 = int(emb_res[0][0] == true_label)
emb_top5 = int(any(lbl == true_label for lbl, _ in emb_res))
emb_dist = emb_res[0][1]
finally:
cur.close()
return raw_top1, raw_top5, raw_dist, emb_top1, emb_top5, emb_dist
def main():
k = 5
n = len(mnist_test)
# 控制并发线程数,避免压垮数据库 (也可用环境变量 OMP_THREADS 指定)
max_workers = int(os.environ.get("QUERY_WORKERS", os.cpu_count() or 4))
# 统计指标
raw_top1_hit = raw_top5_hit = 0
emb_top1_hit = emb_top5_hit = 0
raw_dist_sum = emb_dist_sum = 0.0
with ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = [executor.submit(process_one, idx, k) for idx in range(n)]
for fut in as_completed(futures):
r1, r5, rd, e1, e5, ed = fut.result()
raw_top1_hit += r1
raw_top5_hit += r5
raw_dist_sum += rd
emb_top1_hit += e1
emb_top5_hit += e5
emb_dist_sum += ed
print(f"测试集样本数: {n}, Top-k = {k}")
print("-" * 50)
print("基于原始向量 pixels (欧氏距离 <->):")
print(f" Top-1 命中率: {raw_top1_hit / n:.4f} ({raw_top1_hit}/{n})")
print(f" Top-5 命中率: {raw_top5_hit / n:.4f} ({raw_top5_hit}/{n})")
print(f" 平均最近邻距离: {raw_dist_sum / n:.4f}")
print("-" * 50)
print("基于归一化向量 embedding (欧氏距离 <->):")
print(f" Top-1 命中率: {emb_top1_hit / n:.4f} ({emb_top1_hit}/{n})")
print(f" Top-5 命中率: {emb_top5_hit / n:.4f} ({emb_top5_hit}/{n})")
print(f" 平均最近邻距离: {emb_dist_sum / n:.4f}")
def _close_all_conns():
with _conn_lock:
conns = list(_all_conns)
_all_conns.clear()
for c in conns:
try:
c.close()
except Exception:
pass
if __name__ == "__main__":
try:
main()
finally:
_close_all_conns()
|
