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
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
|
"""
HNSW参数网格搜索与自动调优
帮助找到最优参数组合
"""
import time
import numpy as np
from typing import Dict, List, Tuple
from dataclasses import dataclass
import pandas as pd
import matplotlib.pyplot as plt
@dataclass
class TuningResult:
"""调优结果数据类"""
M: int
efConstruction: int
ef: int
recall: float
latency_ms: float
memory_mb: float
qps: float
score: float
class HNSWParameterOptimizer:
"""HNSW参数自动优化器"""
def __init__(
self,
milvus_client,
collection_name: str,
test_data: Dict,
target_recall: float = 0.95,
max_latency_ms: float = 50.0,
weight_recall: float = 0.4,
weight_speed: float = 0.3,
weight_memory: float = 0.3
):
self.client = milvus_client
self.collection_name = collection_name
self.test_data = test_data
self.target_recall = target_recall
self.max_latency_ms = max_latency_ms
self.weight_recall = weight_recall
self.weight_speed = weight_speed
self.weight_memory = weight_memory
self.results = []
def grid_search(
self,
M_range: List[int] = None,
efConstruction_range: List[int] = None,
ef_range: List[int] = None
) -> pd.DataFrame:
"""
网格搜索最优参数
参数:
M_range: M参数搜索范围
efConstruction_range: efConstruction搜索范围
ef_range: ef搜索范围
"""
if M_range is None:
M_range = [8, 12, 16, 20, 24, 32]
if efConstruction_range is None:
efConstruction_range = [100, 150, 200, 256, 300, 400]
if ef_range is None:
ef_range = [64, 80, 100, 128, 160, 200, 256]
total_combinations = len(M_range) * len(efConstruction_range) * len(ef_range)
print(f"开始网格搜索，共 {total_combinations} 种参数组合...")
completed = 0
for M in M_range:
for efC in efConstruction_range:
for ef in ef_range:
result = self._evaluate_parameters(M, efC, ef)
self.results.append(result)
completed += 1
if completed % 10 == 0:
print(f"进度: {completed}/{total_combinations} ({completed/total_combinations*100:.1f}%)")
results_df = pd.DataFrame([vars(r) for r in self.results])
results_df = results_df.sort_values('score', ascending=False)
return results_df
def _evaluate_parameters(self, M: int, efConstruction: int, ef: int) -> TuningResult:
"""
评估一组参数的性能
"""
try:
start_time = time.time()
self._rebuild_index(M, efConstruction)
recall = self._measure_recall(ef)
latency, qps = self._measure_latency(ef)
memory_mb = self._estimate_memory_usage(M, efConstruction)
elapsed = time.time() - start_time
score = self._calculate_score(recall, latency, memory_mb)
result = TuningResult(
M=M,
efConstruction=efConstruction,
ef=ef,
recall=recall,
latency_ms=latency,
memory_mb=memory_mb,
qps=qps,
score=score
)
return result
except Exception as e:
print(f"评估参数 M={M}, efC={efConstruction}, ef={ef} 时出错: {e}")
return TuningResult(
M=M, efConstruction=efConstruction, ef=ef,
recall=0, latency_ms=999999, memory_mb=999999,
qps=0, score=-1
)
def _rebuild_index(self, M: int, efConstruction: int):
"""重建HNSW索引"""
index_params = self.client.prepare_index_params()
index_params.add_index(
field_name="vector",
index_type="HNSW",
metric_type="COSINE",
params={"M": M, "efConstruction": efConstruction}
)
self.client.drop_index(
collection_name=self.collection_name,
field_name="vector"
)
self.client.create_index(
collection_name=self.collection_name,
index_params=index_params,
sync=True
)
def _measure_recall(self, ef: int, num_queries: int = 100) -> float:
"""测量召回率（与暴力搜索对比）"""
queries = self.test_data['queries'][:num_queries]
ground_truth = self.test_data['ground_truth'][:num_queries]
correct = 0
total = 0
for query_vec, true_neighbors in zip(queries, ground_truth):
results = self.client.search(
collection_name=self.collection_name,
data=[query_vec],
limit=10,
search_params={
"metric_type": "COSINE",
"params": {"ef": ef}
}
)
retrieved_ids = set([hit['id'] for hit in results[0]])
true_ids = set(true_neighbors[:10])
correct += len(retrieved_ids & true_ids)
total += len(true_ids)
recall = correct / total if total > 0 else 0
return recall
def _measure_latency(self, ef: int, num_queries: int = 1000) -> Tuple[float, float]:
"""测量查询延迟和QPS"""
queries = self.test_data['queries'][:num_queries]
latencies = []
start_total = time.time()
for query_vec in queries:
start = time.time()
self.client.search(
collection_name=self.collection_name,
data=[query_vec],
limit=10,
search_params={
"metric_type": "COSINE",
"params": {"ef": ef}
}
)
latency = (time.time() - start) * 1000
latencies.append(latency)
total_time = time.time() - start_total
avg_latency = np.mean(latencies)
p99_latency = np.percentile(latencies, 99)
qps = num_queries / total_time
return p99_latency, qps
def _estimate_memory_usage(self, M: int, efConstruction: int) -> float:
"""估算HNSW索引内存占用（MB）"""
num_vectors = self.test_data.get('num_vectors', 1000000)
dim = 1024
base_memory = num_vectors * dim * 4 / 1024 / 1024
graph_memory = num_vectors * M * 12 / 1024 / 1024
total_memory = base_memory + graph_memory
return total_memory
def _calculate_score(self, recall: float, latency_ms: float, memory_mb: float) -> float:
"""
计算综合评分（加权归一化）
"""
norm_recall = min(recall, 1.0)
norm_speed = 1.0 / (1.0 + np.exp((latency_ms - self.max_latency_ms) / 10))
norm_memory = 1.0 - min(memory_mb / 10240, 1.0)
score = (
self.weight_recall * norm_recall +
self.weight_speed * norm_speed +
self.weight_memory * norm_memory
)
return score
def visualize_results(self, results_df: pd.DataFrame, save_path: str = 'tuning_results.png'):
"""可视化调优结果"""
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
ax1 = axes[0, 0]
scatter = ax1.scatter(
results_df['latency_ms'],
results_df['recall'],
c=results_df['score'],
cmap='viridis',
s=50,
alpha=0.7
)
ax1.set_xlabel('P99 Latency (ms)')
ax1.set_ylabel('Recall')
ax1.set_title('Recall vs Latency (color=Score)')
plt.colorbar(scatter, ax=ax1, label='Score')
ax2 = axes[0, 1]
top_10 = results_df.head(10)
x = np.arange(len(top_10))
width = 0.25
bars1 = ax2.bar(x - width, top_10['M'], width, label='M')
bars2 = ax2.bar(x, top_10['efConstruction'], width, label='efConstruction')
bars3 = ax2.bar(x + width, top_10['ef']/10, width, label='ef/10')
ax2.set_xlabel('Top Configurations')
ax2.set_ylabel('Parameter Value')
ax2.set_title('Top 10 Configurations Parameters')
ax2.legend()
ax2.set_xticks(x)
ax2.set_xticklabels([f'#{i+1}' for i in range(len(top_10))], rotation=45)
ax3 = axes[1, 0]
ax3.scatter(results_df['memory_mb'], results_df['recall'],
c=results_df['score'], cmap='plasma', s=50, alpha=0.7)
ax3.set_xlabel('Memory Usage (MB)')
ax3.set_ylabel('Recall')
ax3.set_title('Memory vs Recall (color=Score)')
ax4 = axes[1, 1]
ax4.hist(results_df['qps'], bins=20, edgecolor='black', alpha=0.7)
ax4.axvline(results_df['qps'].mean(), color='red', linestyle='--', label=f'Mean: {results_df["qps"].mean():.0f}')
ax4.set_xlabel('QPS')
ax4.set_ylabel('Frequency')
ax4.set_title('QPS Distribution')
ax4.legend()
plt.tight_layout()
plt.savefig(save_path, dpi=150, bbox_inches='tight')
plt.close()
print(f"✅ 可视化图表已保存至: {save_path}")
def get_best_configuration(self, results_df: pd.DataFrame) -> Dict:
"""获取最优配置"""
best_row = results_df.iloc[0]
config = {
'M': int(best_row['M']),
'efConstruction': int(best_row['efConstruction']),
'ef': int(best_row['ef']),
'expected_performance': {
'recall': best_row['recall'],
'p99_latency_ms': best_row['latency_ms'],
'memory_mb': best_row['memory_mb'],
'qps': best_row['qps'],
'composite_score': best_row['score']
}
}
return config
if __name__ == "__main__":
"""
from pymilvus import MilvusClient
client = MilvusClient(uri="http://localhost:19530")
# 准备测试数据
test_data = {
'queries': [...], # 测试查询向量列表
'ground_truth': [[...], ...], # 真实最近邻
'num_vectors': 1000000 # Collection中的向量总数
}
optimizer = HNSWParameterOptimizer(
milvus_client=client,
collection_name="rag_vectors",
test_data=test_data,
target_recall=0.95,
max_latency_ms=50.0
)
# 执行网格搜索
results_df = optimizer.grid_search(
M_range=[12, 16, 20],
efConstruction_range=[150, 200, 256],
ef_range=[100, 128, 160]
)
# 可视化结果
optimizer.visualize_results(results_df)
# 获取最优配置
best_config = optimizer.get_best_configuration(results_df)
print("\n=== 最优配置 ===")
print(f"M = {best_config['M']}")
print(f"efConstruction = {best_config['efConstruction']}")
print(f"ef = {best_config['ef']}")
print(f"预期性能: {best_config['expected_performance']}")
"""
pass
|