使用 GPU 建立索引
本指南概述了在 Milvus 中使用 GPU 支援建立索引的步驟,這可以顯著改善高吞吐量和高召回情境下的搜尋效能。有關 Milvus 支援的 GPU 索引類型的詳細資訊,請參閱GPU 索引。
為 GPU 記憶體控制配置 Milvus 設定
Milvus 使用全局繪圖記憶體池來分配 GPU 記憶體。
它在Milvus 配置檔中支援兩個參數initMemSize
和maxMemSize
。記憶體池大 小初始設定為initMemSize
,超過此限制後會自動擴充至maxMemSize
。
當 Milvus 啟動時,預設的initMemSize
是可用 GPU 記憶體的 1/2,而預設的maxMemSize
是等於所有可用的 GPU 記憶體。
在 Milvus 2.4.1 之前(包括 2.4.1 版),Milvus 使用統一的 GPU 記憶池。對於 2.4.1 之前的版本 (包括 2.4.1 版),建議將這兩個值都設定為 0。
gpu:
initMemSize: 0 #set the initial memory pool size.
maxMemSize: 0 #maxMemSize sets the maximum memory usage limit. When the memory usage exceed initMemSize, Milvus will attempt to expand the memory pool.
從 Milvus 2.4.1 起,GPU 記憶體池僅在搜尋時用於臨時 GPU 資料。因此,建議將其設定為 2048 和 4096。
gpu:
initMemSize: 2048 #set the initial memory pool size.
maxMemSize: 4096 #maxMemSize sets the maximum memory usage limit. When the memory usage exceed initMemSize, Milvus will attempt to expand the memory pool.
建立索引
以下範例說明如何建立不同類型的 GPU 索引。
準備索引參數
設定 GPU 索引參數時,請定義index_type、metric_type 和params:
index_type(字串):用於加速向量搜尋的索引類型。有效的選項包括GPU_CAGRA、GPU_IVF_FLAT、GPU_IVF_PQ 及GPU_BRUTE_FORCE。
metric_type(字串):用來衡量向量相似性的度量類型。有效的選項是IP和L2。
params(dict):特定於索引的建立參數。此參數的有效選項取決於索引類型。
以下是不同索引類型的配置範例:
GPU_CAGRA索引
index_params = { "metric_type": "L2", "index_type": "GPU_CAGRA", "params": { 'intermediate_graph_degree': 64, 'graph_degree': 32 } }
params的可能選項包括
intermediate_graph_degree(int):透過在剪枝之前確定圖形的程度來影響召回和建立時間。建議值為32或64。
graph_degree(int):透過設定剪枝後的圖形程度來影響搜尋效能和召回率。通常,它是intermediate_graph_degree 的一半。這兩個程度之間的差異越大,建立時間就越長。它的值必須小於intermediate_graph_degree 的值。
build_algo(字串):選擇剪枝前的圖形生成演算法。可能的選項:
IVF_PQ: 提供較高的品質,但建立時間較慢。
NN_DESCENT:提供較快的建立速度,但召回率可能較低。
cache_dataset_on_device(string,"true" | "false"):決定是否在 GPU 記憶體中快取原始資料集。將此設定為"true 「可以精煉搜尋結果,從而提高召回率,而設定為」false "則可節省 GPU 記憶體。
GPU_IVF_FLAT或GPU_IVF_PQ索引
index_params = { "metric_type": "L2", "index_type": "GPU_IVF_FLAT", # Or GPU_IVF_PQ "params": { "nlist": 1024 } }
GPU_BRUTE_FORCE索引
index_params = { 'index_type': 'GPU_BRUTE_FORCE', 'metric_type': 'L2', 'params': {} }
不需要額外的params設定。
建立索引
在index_params 中設定索引參數後,呼叫 create_index()
方法來建立索引。
# Get an existing collection
collection = Collection("YOUR_COLLECTION_NAME")
collection.create_index(
field_name="vector", # Name of the vector field on which an index is built
index_params=index_params
)
搜尋
建立 GPU 索引後,下一步就是在進行搜尋前準備搜尋參數。
準備搜尋參數
以下是不同索引類型的配置範例:
GPU_BRUTE_FORCE索引
search_params = { "metric_type": "L2", "params": {} }
不需要額外的參數配置。
GPU_CAGRA索引
search_params = { "metric_type": "L2", "params": { "itopk_size": 128, "search_width": 4, "min_iterations": 0, "max_iterations": 0, "team_size": 0 } }
主要搜尋參數包括
itopk_size:決定搜尋過程中保留的中間結果大小。較大的值可能會提高召回率,但卻會犧牲搜尋效能。它應該至少等於最終 top-k(限制) 值,通常是 2 的幂次 (例如 16、32、64、128)。
search_width:指定搜尋期間進入 CAGRA 圖形的入口點數量。增加此值可提高召回率,但可能會影響搜尋效能。
min_iterations/max_ iterations:這些參數控制搜尋的迭代程序。預設值為0,CAGRA 會根據itopk_size和search_width 自動決定迭代次數。手動調整這些值有助於平衡效能與精確度。
team_size:指定用於計算 GPU 公制距離的 CUDA 線程數目。常見的值是 2 的幂數,最高為 32 (例如 2、4、8、16、32)。它對搜尋效能影響不大。預設值是0,Milvus 會根據向量的維度自動選擇team_size。
GPU_IVF_FLAT或GPU_IVF_PQ索引
search_params = { "metric_type": "L2", "params": {"nprobe": 10} }
進行搜尋
使用 search()
方法在 GPU 索引上執行向量相似性搜尋。
# Load data into memory
collection.load()
collection.search(
data=[[query_vector]], # Your query vector
anns_field="vector", # Name of the vector field
param=search_params,
limit=100 # Number of the results to return
)
限制
使用 GPU 索引時,請注意某些限制:
對於GPU_IVF_FLAT,limit的最大值為 1024。
對於GPU_IVF_PQ和GPU_CAGRA,limit的最大值為 1024。
雖然GPU_BRUTE_FORCE 沒有設定限制,但建議不要超過 4096,以避免潛在的效能問題。
目前,GPU 索引不支援 COSINE 距離。如果需要 COSINE 距離,應該先將資料規格化,然後再使用內積 (IP) 距離來替代。
不完全支援 GPU 索引的載入 OOM 保護,太多資料可能會導致 QueryNode 當機。
常見問題
何時適合使用 GPU 索引?
GPU 索引特別適用於需要高吞吐量或高召回率的情況。例如,在處理大量批次時,GPU 索引的吞吐量可比 CPU 索引高出 100 倍之多。在批次較小的情況下,GPU 索引的效能仍遠遠超越 CPU 索引。此外,如果需要快速插入資料,整合 GPU 可大幅加快建立索引的過程。
GPU 索引(如 CAGRA、GPU_IVF_PQ、GPU_IVF_FLAT 和 GPU_BRUTE_FORCE)最適合哪些應用場合?
CAGRA 索引非常適合需要增強效能的應用環境,儘管其代價是消耗更多的記憶體。對於以節省記憶體為優先考量的環境,GPU_IVF_PQ索引可幫助將儲存需求降至最低,儘管這會帶來較高的精確度損失。GPU_IVF_FLAT索引是一個平衡的選擇,提供效能與記憶體使用量之間的折衷方案。最後,GPU_BRUTE_FORCE索引專為窮盡搜尋作業而設計,透過執行遍歷搜尋來保證召回率為 1。