8.8 Image Retrieval 圖像檢索 （下）
前言
上一小節，對圖像檢索的基礎概念進行了介紹，本節將通過代碼實踐，實現以文搜圖、以圖搜圖的功能。
本節代碼核心包括：
- Faiss 框架介紹及常用演算法評估，並基於Faiss實現COCO 2017的11萬資料的圖像檢索
- CLIP實現image/text的特徵提取
- 集成Faiss+CLIP構建無需訓練的圖像檢索系統
- 基於Flask將圖像檢索系統部署為web服務
Faiss安裝及使用
簡介
Faiss是MetaAI開源的高性能向量檢索庫，它基於C++編寫，提供python介面，核心模組還可以用GPU加速，在億級數據可在xxx秒級實現結果返回，目前應用較為廣泛。中小型項目及個人，非常適合採用Faiss。
Faiss目前最好的學習資源有兩個，一個是官方wiki，一個是www.pinecone.io的Faiss: The Missing Manual
faiss提供的優化演算法主要分為PQ量化和倒排索引，下圖為faiss中演算法體系結構圖，可以看到核心在藍色區域和黃色區域。
安裝
可以通過pip或者conda安裝，這裡推薦conda，這裡需要注意鏡像源最好conda源，若是清華鏡像是行不通的。
# cpu版本安裝
conda install -c pytorch faiss-cpu
# gpu版本安裝
conda install -c pytorch faiss-gpu
Copy
切換回到conda官方源方法：
# 依次輸入
conda config --remove-key channels
conda config --add channels defaults
conda config --add channels conda-forge
Copy
在這裡安裝了gpu版本，便於測試，在gpu版本安裝中，會默認裝上cudatoolkit， 600多M，總共需要900多M
The following packages will be downloaded:
package | build
---------------------------|-----------------
ca-certificates-2022.12.7 | h5b45459_0 143 KB conda-forge
cudatoolkit-11.8.0 | h09e9e62_11 638.9 MB conda-forge
faiss-1.7.2 |py38cuda112h7f1466e_3_cuda 885 KB conda-forge
faiss-gpu-1.7.2 | h949689a_3 15 KB conda-forge
intel-openmp-2023.1.0 | h57928b3_46319 2.5 MB conda-forge
libblas-3.9.0 | 16_win64_mkl 5.6 MB conda-forge
libcblas-3.9.0 | 16_win64_mkl 5.6 MB conda-forge
libfaiss-1.7.2 |cuda112h33bf9e0_3_cuda 51.0 MB conda-forge
libfaiss-avx2-1.7.2 |cuda112h1234567_3_cuda 51.1 MB conda-forge
libhwloc-2.9.1 | h51c2c0f_0 2.4 MB conda-forge
libiconv-1.17 | h8ffe710_0 698 KB conda-forge
liblapack-3.9.0 | 16_win64_mkl 5.6 MB conda-forge
libxml2-2.10.4 | hc3477c8_0 1.7 MB conda-forge
libzlib-1.2.13 | hcfcfb64_4 70 KB conda-forge
mkl-2022.1.0 | h6a75c08_874 182.7 MB conda-forge
numpy-1.24.2 | py38h7ec9225_0 5.6 MB conda-forge
openssl-1.1.1t | hcfcfb64_0 5.0 MB conda-forge
pthreads-win32-2.9.1 | hfa6e2cd_3 141 KB conda-forge
python_abi-3.8 | 2_cp38 4 KB conda-forge
setuptools-67.6.1 | pyhd8ed1ab_0 567 KB conda-forge
tbb-2021.9.0 | h91493d7_0 151 KB conda-forge
ucrt-10.0.22621.0 | h57928b3_0 1.2 MB conda-forge
vs2015_runtime-14.34.31931 | h4c5c07a_10 708 KB conda-forge
------------------------------------------------------------
Total: 962.3 MB
Copy
安裝驗證——"Hello Faiss"
接下來採用faiss進行精確查找，實現10000條查詢向量的top-4向量檢索。
faiss使用步驟主要分三步：
- 創建索引子，索引子有FlatL2、LSH、PQ、HNSWFlat等
- 初始化索引子，將資料庫向量添加到索引子中，並進行預訓練（如果需要）
- 使用索引子進行檢索。
以下為配套代碼運行結果
import time
import numpy as np
import faiss
# ============================ step 0: 數據構建 ============================
np.random.seed(1234)
d = 64 # dimension
nb = 100000 # database size
nq = 10000 # nb of queries
xb = np.random.random((nb, d)).astype('float32')
xq = np.random.random((nq, d)).astype('float32')
xb[:, 0] += np.arange(nb) / 1000.
xq[:, 0] += np.arange(nq) / 1000.
# ============================ step 1: 構建索引子 ============================
index = faiss.IndexFlatL2(d)
index.add(xb)
# ============================ step 2: 索引 ============================
k = 4 # top_k number
for i in range(5):
s = time.time()
D, I = index.search(xq, k)
print("{}*{}量級的精確檢索，耗時:{:.3f}s".format(nb, nq, time.time()-s))
# ============================ step 3: 檢查索引結果 ============================
print('D.shape: {}, D[0, ...]: {}'.format(D.shape, D[0]))
print('I.shape: {}, I[0, ...]: {}'.format(I.shape, I[0]))
# D是查詢向量與topk向量的距離，distance
# I是與查詢向量最近的向量的id，此處有10萬資料，index在0-99999之間。
Copy
輸出的D表示距離， 6.815表示第一個查詢向量的top-1向量的距離是6.815
輸出的I表示向量id， 381表示第一個查詢向量的top-1向量的id是381
D.shape: (10000, 4), D[0, ...]: [6.815506 6.8894653 7.3956795 7.4290257]
I.shape: (10000, 4), I[0, ...]: [381 207 210 477]
Copy
faiss提供的方法匯總表如下：
|
Method |
Class name |
index_factory |
Main parameters |
Bytes/vector |
Exhaustive |
Comments |
|---|---|---|---|---|---|---|
|
Exact Search for L2 |
IndexFlatL2 |
"Flat" |
d |
4*d |
yes |
brute-force |
|
Exact Search for Inner Product |
IndexFlatIP |
"Flat" |
d |
4*d |
yes |
also for cosine (normalize vectors beforehand) |
|
Hierarchical Navigable Small World graph exploration |
IndexHNSWFlat |
'HNSWx,Flat` |
d,M |
4d + x M 2 4 |
no |
|
|
Inverted file with exact post-verification |
IndexIVFFlat |
"IVFx,Flat" |
quantizer,d,nlists,metric |
4*d + 8 |
no |
Takes another index to assign vectors to inverted lists. The 8 additional bytes are the vector id that needs to be stored. |
|
Locality-Sensitive Hashing (binary flat index) |
IndexLSH |
- |
d,nbits |
ceil(nbits/8) |
yes |
optimized by using random rotation instead of random projections |
|
Scalar quantizer (SQ) in flat mode |
IndexScalarQuantizer |
"SQ8" |
d |
d |
yes |
4 and 6 bits per component are also implemented. |
|
Product quantizer (PQ) in flat mode |
IndexPQ |
"PQx","PQ"x"x"nbits |
d,M,nbits |
ceil(M * nbit / 8) |
yes |
|
|
IVF and scalar quantizer |
IndexIVFScalarQuantizer |
"IVFx,SQ4" "IVFx,SQ8" |
quantizer,d,nlists,qtype |
SQfp16: 2 *d+ 8, SQ8:d+ 8 or SQ4:d/2+ 8 |
no |
Same as theIndexScalarQuantizer |
|
IVFADC (coarse quantizer+PQ on residuals) |
IndexIVFPQ |
"IVFx,PQ"y"x"nbits |
quantizer,d,nlists,M,nbits |
ceil(M * nbits/8)+8 |
no |
|
|
IVFADC+R (same as IVFADC with re-ranking based on codes) |
IndexIVFPQR |
"IVFx,PQy+z" |
quantizer,d,nlists,M,nbits,M_refine,nbits_refine |
M+M_refine+8 |
no |
faiss 基礎
相似性評價指標
faiss的評價指標主要有L2 和 inner product，L2是平方後的L2，這是為了減少計算量，只是比較大小，就沒必要開平方了，需要真正意義的L2，要自行開平方。
除了L2和inner product，還有METRIC_L1, METRIC_Linf and METRIC_Lp ，但不常用，需要時查看官方wiki文檔。
faiss的數據預處理
faiss提供了高性能的 k-means clustering, PCA, PQ encoding/decoding，需要用時查看wiki
索引子太多，如何選？
- RAM充足：選HNSW，IVF1024,PQNx4fs,RFlat
- RAM一般：OPQ M _D ,...,PQ M x 4fsr
- RAM不足：OPQM _D ,...,PQM
- 資料量不大，且要求精確：選 IndexFlatL2 或 IndexFlatIP
- 數據量大：IVF K, K可以選擇256，2^16=65536， 2^18=262144， 2^20=1048576，根據資料量在1M, 10M, 100M, 1B之間選擇。
綜上：記憶體問題用量化，速度問題用倒排，一個普適的方法是 IVF + PQ。
faiss預處理及後處理
預處理方法：
- random rotation， RandomRotationMatrix， useful to re-balance components of a vector before indexing in an IndexPQ or IndexLSH
- remapping of dimensions， RemapDimensionsTransform， to reduce or increase the size of a vector because the index has a preferred dimension, or to apply a random permutation on dimensions.
- PCA, PCAMatrix, for dimensionality reduction ,
- OPQ rotation, OPQMatrix, OPQ applies a rotation to the input vectors to make them more amenable to PQ coding. See Optimized product quantization, Ge et al., CVPR'13 for more details.
後處理方法：
- re-ranking：IndexRefineFlat，基於距離的重排，利用近似檢索獲得的結果精度可能較低，可通過距離進行重排。
- IndexShards：組合多個索引子的結果，或是多gpu並行結果的融合。
index factory
faiss提供了基於字串創建索引子的工廠函數，字串以逗號分隔，可以一次性創建複合的索引子，包括預處理、索引、後處理等。
舉例：
- index = index_factory(128, "PCA80,Flat") ：為 128維 向量生成一個索引，通過 PCA 將它們減少到 80維，然後進行精確索引。
- index = index_factory(128, "OPQ16_64,IMI2x8,PQ8+16")：輸入是128維度 向量，再將 OPQ 變換應用於 64D 中的 16 個塊，再使用 2x8 位元的反向多索引（= 65536 個反向列表），最後使用大小為 8 的 PQ 進行量化16 位元組。
索引資料的I/O
資料庫的索引資訊通常需要存於磁片，再讀入記憶體，這時候需要讀取I/O方法。
write_index(index, "large.index"): writes the given index to file large.index
index = read_index("large.index"): reads a file
Copy
GPU 使用
faiss提供核心索引子的gpu加速，可以將資料放到gpu上進行加速運算，使用比較方便，只需要以下三步：先獲取gpu資源、創建cpu的索引子、cpu索引子搬到gpu上。
注意：PQ不支援gpu加速，IVFPQ才支援。GPU: only pq.nbits == 8 is supported
res = faiss.StandardGpuResources() # 1. 獲取gpu資源
index_flat = faiss.IndexFlatL2(d) # 2. 創建cpu索引子
gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index_flat) # 3. 遷移至gpu
Copy
索引速度比cpu快5-10倍，筆記本測試是快了10倍，詳情可運行配套代碼。
faiss還提供多gpu並行運算，多gpu使用如下代碼所示
import numpy as np
import faiss
d = 64 # dimension
nb = 100000 # database size
nq = 10000 # nb of queries
np.random.seed(1234) # make reproducible
xb = np.random.random((nb, d)).astype('float32')
xb[:, 0] += np.arange(nb) / 1000.
xq = np.random.random((nq, d)).astype('float32')
xq[:, 0] += np.arange(nq) / 1000.
ngpus = faiss.get_num_gpus()
print("number of GPUs:", ngpus)
cpu_index = faiss.IndexFlatL2(d)
gpu_index = faiss.index_cpu_to_all_gpus(cpu_index)
gpu_index.add(xb) # add vectors to the index
print(gpu_index.ntotal)
k = 4 # we want to see 4 nearest neighbors
D, I = gpu_index.search(xq, k) # actual search
print('D.shape: {}, D[0, ...]: {}'.format(D.shape, D[0]))
print('I.shape: {}, I[0, ...]: {}'.format(I.shape, I[0]))
Copy
除了cpu遷移至gpu和多gpu的函數，還有gpu遷移至cpu的函數：faiss.index_gpu_to_cpu。
使用gpu需要注意：
- k and nprobe must be <= 2048 for all indices.
- For GpuIndexIVFPQ, code sizes per encoded vector allowed are 1, 2, 3, 4, 8, 12, 16, 20, 24, 28, 32, 48, 56, 64 and 96 bytes.
- 其它：詳見wiki
- cpu與gpu有相同的精度，但是會發生cpu與gpu的檢索結果不一致的情況！這可能是floating point reduction order、equivalent element k-selection order、float16 opt-in導致的
faiss 代碼結構設計
faiss 底層代碼為CUDA、BLAS、numpy，其中CUDA+BAL是核心，構建了所有底層計算（淺綠色方框），再往上就是對python、Rust/C#、C++的代碼封裝，用戶可直接使用python等語言調用faiss的功能。
這裡主要看python部分，對於底層計算庫c++/cuda，經過SWIG將numpy與底層計算庫封裝為python類，變為python可調用的介面；向上再經過Adaptor layer封裝，以便於python物件可轉換為C++物件；隨後採用contrib library進行封裝，提供python代碼介面。
- SWIG (Simplified Wrapper and Interface Generator) 是一個用於將 C/C++ 代碼封裝為其他程式設計語言介面的開源工具。
- Adaptor layer 是指用 Python 編寫的橋接層，其作用是將 Python 物件與 C++ 物件進行轉換，通過定義 Python 類或函數的方式，將 C++ 類或函數封裝為 Python 物件或函數
- Python contrib library 是一個開源的 Python 庫，它包含了許多常用的機器學習和深度學習模型、工具和資料集。
faiss有一大特點，可以用pytorch的tensor可直接傳入索引子的.add() 和 .serarch()。
faiss進階知識點
- 執行緒與非同步：cpu是執行緒安全的，gpu不是執行緒安全的，詳情參見wiki
- 底層API介紹：InvertedLists 和 InvertedListsScannerwiki
- 超大資料的存儲方案：分散式存儲於多機器、存儲於磁片、分散式鍵值存儲wiki
- 二進位量化編碼：三個核心API及使用demowiki
- 暴力檢索的直接調用：暴力檢索可以不需要索引子，直接調用函數計算就好。wiki
- 低比特高速索引：借鑒Google的SCANN 4-bit PQ fast scan，實現高速索引。wiki)
- 實戰筆記系列：k-means演算法、IVFPQ的預計算表、PCA矩陣計算、非窮舉搜索的統計資訊、重排refine的作用及參數設置。wiki
- 超參數自動搜索：例如nprobe，有的向量用10，有的向量用20，這裡提供了SearchParameter實現自動搜索。wiki
- 加速優化策略：大記憶體頁、選擇適當的CPU資源、選擇適當的MKL執行緒數等。wiki
Faiss 的benchmark
接下來對常用的PQ與IVFPQ進行benchmark測試，觀察不同參數下，耗時與recall的情況，為後續實際應用挑選參數做準備。
資料集下載：採用sift1M資料（http://corpus-texmex.irisa.fr/），提供100萬的圖像特徵向量，向量長度為128，在個人電腦上可運行。
首先觀察Product Quantization演算法的參數，PQ中子段的數量會影響計算速度與召回，子段越多召回越高，同時量化bit越小，記憶體佔用越小，但召回降低。
下面就觀察子段分別是4,8,16,32時，量化bit分別是6,7,8,9時，耗時與召回的情況。
運行配套代碼，可得到如下兩張圖，分別是召回的對比、耗時的對比。
通過實驗結果圖可知道：
- 16個子段大約是個分水嶺，小於16個子段時，召回會快速提升。一般用16個子段、32個子段
- 耗時方面，8bit的耗時出現異常，這可能是由於代碼針對8bit做了特殊優化，並且論文中推薦的也是8bit量化
- 後續使用，直接用PQ16x8或者PQ32x8就好。
PQ + IVF
在faiss中PQ通常不會單獨使用，而是結合IVF，IVF中聚類中心的數量會影響速度和召回，檢索時probe的數量也會影響速度和召回。
由於PQ參數已經選好了，這裡採用PQ32x8，IVF的聚類中心分別有128，256，512，1024，2048，4096；probe分別有1, 2, 4, 8, 16, 32, 64。
運行配套代碼，可得到如下兩張圖，分別是召回的對比、耗時的對比。
通過實驗結果圖可知道：
- 相同probe時，聚類中心越少，召回越高。可能是資料集的問題，聚類中心的差異沒有很明顯，需要把128和4096進行對比，精度差別不大。
- 聚類中心越多，耗時越少，因為一個類裡邊的向量個數少了，需要匹配的向量就少了，速度就快
- probe數量增大，耗時增加，在32之前增加較小，因此32是個不錯的選擇，當精度不足時可適當增加到48。PQ的wiki中也提到IVFPQ的對比，probe=48時，可以得到與PQ相同水準的recall，並且耗時大幅度較少。
- 聚類中心選4096，nprobe選48，下調最多到36，否則精度降低，並且速度沒啥提升。
經過兩個實驗，可以得到一個通用的檢索器，"IVF4096,PQ32x8"，並且推理檢索時，設置index.nprobe=48。
更多benchmark及應用案例可以看wiki中的Case-studies, benchs
基於CLIP+Faiss+Flask的圖像檢索系統
接下來，採用CLIP模型充當特徵提取器，採用Faiss實現向量檢索，採用Flask進行web服務部署，最終可實現以圖搜圖和以文搜圖。
回顧圖像檢索簡介中介紹的，一個圖像檢索系統如下圖所示，
在構建圖像檢索系統時，需要關注：
- 特徵提取器：採用CLIP(Contrastive Language-Image Pre-training)模型進行圖像特徵提取，是一個文圖預訓練模型，CLIP於2021年2月由openAI發表，並開源了模型，模型由4億的圖文資料，採用對比學習方式進行訓練得到，由於對比學習與超大規模的資料集加持，使CLIP模型很好的理解了自然圖像，在眾多資料集上表現出了優異的zero-shot性能，同時在表徵學習（representation learning）中也很好。CLIP簡介可以回顧圖像描述，上一節採用了CLIP進行圖像描述。
- 向量檢索：基於faiss的IVF+PQ方法
- 推薦策略：目前缺少業務背景及需求，這裡不涉及。
代碼結構
整套代碼位於：pytorch-tutorial-2nd\code\chapter-8\08_image_retrieval，代碼結構如下：
├── config
├── data
├── flask_app.py
├── image_feature_extract.py
├── my_utils
├── retrieval_by_faiss.py
├── static
└── templates
Copy
需要運行的為三份.py腳本分別是
- image_feature_extract.py：特徵向量庫構建，基於CLIP將資料圖片進行特徵提取，並存儲為特徵向量形式
- retrieval_by_faiss.py：基於faiss+clip的圖像檢索演算法實現，作為被調用的代碼塊
- flask_app.py：web服務部署，將調用retrieval_by_faiss中的ImageRetrievalModule實現檢索
其它資料夾功能如下：
- config：整個專案的設定檔，需要設置圖片資料根目錄、faiss檢索演算法超參、相似topk等
- data：存放clip模型提取的特徵，即特徵資料庫。該路徑可以在設定檔中修改
- my_utils：功能函數
- static：flask框架靜態檔的目錄，很重要的一點是需要在裡面構建軟連結，讓flask能訪問到coco2017的圖片，否則無法在web前端頁面展示
- templates：flask的範本目錄，存放html檔，前端頁面設計就在index.htm中
整個圖像檢索運行邏輯如下圖所示：
準備階段，用clip對資料庫進行資料編碼，得到圖像特徵字典庫，將特徵向量傳遞給faiss，用於構建檢索器；圖片id及路徑關係用於id到路徑的映射
推理階段，圖片或文本經clip編碼，輸入faiss進行向量檢索，得到向量id，再由{id:path}字典，獲得圖片路徑，最後進行視覺化。
PS：這裡有一個需要注意的是，文本特徵向量與圖像特徵向量進行匹配時，是需要進行norm操作的！在這裡被坑了一天，最好對照論文才發現少了norm這個步驟。
代碼UML設計簡圖如下：
將特徵提取與向量檢索分別設計一個類來實現，最後將它們這到圖像檢索模組中，對外提供圖像檢索服務。
代碼使用
第一步，配置路徑：修改config檔，設置圖片檔路徑等資訊
第二步，進行特徵提取：運行image_feature_extract.py，進行特徵提取，在data資料夾下獲得兩個pkl文件
第三步，啟動web服務，運行flask_app.py，可到http://localhost:5000/，使用圖像檢索系統
圖像檢索系統web使用說明
PS：由於前端知識的欠缺，有意美化該介面的，歡迎提issue！
小結
本小節通過代碼介紹了faiss庫的基礎使用，並對常見的PQ， IVF+PQ的性能進行了評測，對未來使用向量檢索框架做準備；
同時還涉及了基於CLIP+Faiss+Flask的圖像檢索系統部署，其中使用了優秀的多模態模型——CLIP進行特徵提取，通過這一套系統可以鞏固圖像檢索系統構建時的幾個要點：預處理建特徵向量庫；檢索器初始化；線上檢索服務。
圖像檢索領域涵蓋的知識點太多了，本案例僅能作為入門教程，各細分方向則需要大家自行修煉，祝好！
留言列表
