TiDB Vector 搶先體驗之用 TiDB 實現以圖搜圖

balahoho發表於2024-04-23
TiDB

本文首發自 TiDB 社群專欄：https://tidb.net/blog/0c5672b9
轉載請註明出處！

前言

最早知道 TiDB 要支援向量化的訊息應該是在23年10月份左右，到第一次見到 TiDB Vector 的樣子是在今年1月初，當時 dongxu 在朋友圈發了一張圖：

Weixin Image_20240416150834.jpg

去年我研究了一段時間的向量資料庫，一直對 TiDB 向量特性非常期待，看到這張圖真的就激動萬分，於是第一時間提交了 waitlist 等待體驗 private beta。

苦等幾個月，它終於來了（目前只對 TiDB Serverless 開放）。迫不及待做個小應用嚐嚐鮮。

waitlist申請入口：https://tidb.cloud/ai

體驗入口：https://tidbcloud.com/

建立 TiDB Vector 例項

在收到體驗邀請郵件後，恭喜你可以開始 TiDB Vector 之旅了。

TiDB Serverless 提供了免費試用額度，對於測試用途綽綽有餘，只需要註冊一個 TiDB Cloud 賬號即可。

建立 TiDB Vector 例項和普通的 TiDB 例項並沒有太大區別，在建立叢集頁面可以看到加入瞭如下開關：

企業微信截圖_20240412150805.png

不過要注意的是目前 TiDB Vector 只在 AWS 的eu-central-1可用區開放，選到了其他可用區就看不到這個開關。

這裡只需要填一個叢集名稱就可以開始建立，建立成功後的樣子如下所示：

企業微信截圖_20240412150930.png

下面開始進入正題。

關於向量的那些事

一些基礎概念

  • 向量：向量就是一組浮點數，在程式語言中通常體現為 float 陣列，陣列的長度叫做維度（dim），維度越大精度越高，向量的數學表示是多維座標系中的一個點。例如RGB顏色表示法就是一個簡單的向量示例。
  • embedding：中文翻譯叫嵌入，感覺不好理解，實質上就是把非結構化資料（文字、語音、圖片、影片等）透過一系列演算法加工變成向量的過程，這裡面的演算法叫做模型（model）。
  • 向量檢索：計算兩個向量之間的相似度。

向量檢索初體驗

連線到 TiDB Serverless 後，就可以體驗文章開頭圖片中的向量操作。

建立一張帶有向量欄位的表，長度是3維。

CREATE TABLE vector_table (
    id int PRIMARY KEY,
    doc TEXT,
    embedding vector < float > (3)
  );

往表中插入向量資料：

INSERT INTO vector_table VALUES (1, 'apple', '[1,1,1]'), (2, 'banana', '[1,1,2]'), (3, 'dog', '[2,2,2]');

根據指定的向量做搜尋：

SELECT *, vec_cosine_distance(embedding, '[1,1,3]') as distance FROM vector_table ORDER BY distance LIMIT 3;

+-----------------------+-----------------------+---------------------+
| id      | doc         | embedding             | distance            |
+-----------------------+-----------------------+---------------------+
| 2       | banana      | [1,1,2]               | 0.015268072165338209|
| 3       | dog         | [2,2,2]               | 0.1296117202215108  |
| 1       | apple       | [1,1,1]               | 0.1296117202215108  |
+---------+-------------+-----------------------+---------------------+

這裡的distance就是兩個向量之間的相似度，這個相似度是用vec_cosine_distance函式計算出來的，意味著兩個向量之間的夾角越小相似性越高，夾角大小用餘弦值來衡量。

還有以一種常用的相似度計算方法是比較兩個向量之間的直線距離，稱為歐式距離。

這也意味著不管兩個向量是否有關聯性，總是能計算出一個相似度，distance越小相似度越高。

向量檢索原理

前面大概也提到了兩種常用的向量檢索方式：餘弦相似度和歐式距離，不妨從從最簡單的二維向量開始推導一下計算過程。

二維向量對應一個平面座標系，一個向量就是座標系中任意一點，要計算兩點之間的直線距離用勾股定理很容易就能得出，兩點夾角的餘弦值也有公式能直接算出來。

擴充到三維座標系，還是套用上一步的數學公式，只是多了一個座標。

以此類推到n維也是一樣的方法。

640.webp

640 (1).webp

以上內容來自我去年講的向量資料庫公開課：https://www.bilibili.com/video/BV1YP411t7Do

可以發現維數越多，對算力的要求就越高，計算時間就越長。

第一個 TiDB AI 應用：以圖搜圖

基礎實現

藉助前面介紹的理論知識，一個以圖搜圖的流程應該是這樣子：

企業微信截圖_20240417113526.png

下面我用最簡潔直白的程式碼演示整個流程，方便大家理解。

首先肯定是先連線到 TiDB 例項，目前官方提供了python SDK包tidb_vector，對SQLAlchemy、Peewee這樣的 ORM 框架也有支援，具體可參考https://github.com/pingcap/tidb-vector-python

這裡簡單起見直接用pymysql手寫 SQL 操作，以下連線引數都可以從 TiDB Cloud 控制檯獲取：

import pymysql

def GetConnection():
    connection = pymysql.connect(
        host = "xxx.xxx.prod.aws.tidbcloud.com",
        port = 4000,
        user = "xxx.root",
        password = "xxx",
        database = "test",
        ssl_verify_cert = True,
        ssl_verify_identity = True,
        ssl_ca = "C:\\Users\\59131\\Downloads\\isrgrootx1.pem"
    )
    return connection

再借助 Towhee 來簡化 embedding 的處理，裡面包含了常用的非結構化資料到向量資料的轉換模型，用流水線（pipeline）的形式清晰構建整個處理過程。

from towhee import ops,pipe,AutoPipes,AutoConfig,DataCollection

image_pipe = AutoPipes.pipeline('text_image_embedding')

這裡使用預設配置構建了一個text_image_embedding流水線，它專門用於對文字和圖片做向量轉換，從引用的原始碼中可以看到它使用的模型是clip_vit_base_patch16，預設模態是image。

@AutoConfig.register
class TextImageEmbeddingConfig(BaseModel):
    model: Optional[str] = 'clip_vit_base_patch16'
    modality: Optional[str] = 'image'
    customize_embedding_op: Optional[Any] = None
    normalize_vec: Optional[bool] = True
    device: Optional[int] = -1

clip_vit_base_patch16是一個512維的模型，因此需要在 TiDB 中建立512維的向量欄位。

create table if not exists img_list 
(
    id int PRIMARY KEY, 
    path varchar(200) not null, 
    embedding vector<float>(512)
);

我準備了3000張各種各樣的動物圖片用於測試，把它們依次載入到 TiDB 中，完整程式碼為：

def LoadImage(connection):
    cursor = connection.cursor() 
    cursor.execute("create table if not exists img_list (id int PRIMARY KEY, path varchar(200) not null, embedding vector<float>(512));")
    img_dir='D:\\\\test\\\\'
    files = os.listdir(img_dir)
    for i in range(len(files)):
        path=os.path.join(img_dir, files[i])
        embedding = image_pipe(path).get()[0]
        cursor.execute("INSERT INTO img_list VALUE ("+str(i)+",'"+path+"' , '"+np.array2string(embedding, separator=',')+"');")
    connection.commit()

如果用 ORM 框架的話這裡對資料庫和向量加工操作會簡單些，不需要陣列到字串之間的手工轉換。

載入完成後的資料：

企業微信截圖_20240416181205.png

企業微信截圖_20240417100300.png

下一步定義出根據指定向量在 TiDB 中檢索的函式：

def SearchInTiDB(connection,vector):
    cursor = connection.cursor() 
    begin_time = datetime.datetime.now()
    cursor.execute("select id,path,vec_cosine_distance(embedding, '"+np.array2string(vector, separator=',')+"') as distance from img_list order by distance limit 3;")
    end_time=datetime.datetime.now()
    print("Search time:",(end_time-begin_time).total_seconds())
    df =pd.DataFrame(cursor.fetchall())
    return df[1]

這裡根據餘弦相似度取出結果最相近的3張圖片，返回它們的檔案路徑用於預覽顯示。

下一步用相同的 image pipeline 給指定圖片做 embedding 得到向量，把這個向量傳到 TiDB 中去搜尋，最後把搜尋結果輸出顯示。

def read_images(img_paths):
    imgs = []
    op = ops.image_decode.cv2_rgb()
    for p in img_paths:
        imgs.append(op(p))
    return imgs
    
def ImageSearch(connection,path):    
    emb = image_pipe(path).get()[0]
    res = SearchInTiDB(connection,emb)
    p = (
        pipe.input('path','search_result')
        .map('path', 'img', ops.image_decode.cv2('rgb'))
        .map('search_result','prev',read_images)
        .output('img','prev')
    )
    DataCollection(p(path,res)).show()

看一下最終搜尋效果如何。先看一張已經在圖片庫存在的圖（左邊是待搜尋的圖，右邊是搜尋結果，按相似度由高到低）：

企業微信截圖_20240417101029.png

不能說非常相似，只能說是一模一樣，準確度非常高！再看一下不在圖片庫的搜尋效果：

企業微信截圖_20240417102043.png

企業微信截圖_20240417102455.png

圖片庫裡有幾十種動物，能夠準確搜尋出需要的是狗，特別是第一張從圖片色彩、畫面角度、動作神態上來說都非常相似。

使用向量索引最佳化

沒錯，向量也能加索引，但這個索引和傳統的 B+ Tree 索引有些區別。前面提到向量相似度計算是一個非常消耗 CPU 的過程，如果每次計算都採用全量暴力搜尋的方式那麼無疑效率非常低。上一節演示的案例就是用指定的向量與表裡的3000個向量逐一計算，最簡單粗暴的辦法。

向量索引犧牲了一定的準確度來提升效能，通常採用 ANN（近似最近鄰搜尋） 演算法，HNSW 是最知名的演算法之一。TiDB Vector 目前對它已經有了支援：

create table if not exists img_list_hnsw 
(
    id int PRIMARY KEY, 
    path varchar(200) not null, 
    embedding vector<float>(512) COMMENT "hnsw(distance=cosine)"
);

重新把3000張圖片載入到新的img_list_hnsw表做搜尋測試。

以下分別是不帶索引和帶索引的查詢耗時，第二次明顯要快很多，如果資料量越大這個差距會越明顯，只是目前還無法透過執行計劃或其他方式區分出索引使用情況。

E:\GitLocal\AITester>python tidb_vec.py
Search time: 0.320241
+------------------------------------+------------------------------------------------------------------------------------------------------+
| img                                | prev                                                                                                 |
+====================================+======================================================================================================+
| Image shape=(900, 900, 3) mode=RGB | [Image shape=(84, 84, 3) mode=RGB,Image shape=(84, 84, 3) mode=RGB,Image shape=(84, 84, 3) mode=RGB] |
+------------------------------------+------------------------------------------------------------------------------------------------------+

E:\GitLocal\AITester>python tidb_vec.py
Search time: 0.239746
+------------------------------------+------------------------------------------------------------------------------------------------------+
| img                                | prev                                                                                                 |
+====================================+======================================================================================================+
| Image shape=(900, 900, 3) mode=RGB | [Image shape=(84, 84, 3) mode=RGB,Image shape=(84, 84, 3) mode=RGB,Image shape=(84, 84, 3) mode=RGB] |
+------------------------------------+------------------------------------------------------------------------------------------------------+

實際在本次測試中發現，使用 HNSW 索引對搜尋結果準確度沒有任何影響。

自然語言實現圖片搜尋

本來到這裡測試目的已經達到了，突發奇想想試一下用自然語言也來實現圖片搜尋。於是對程式碼稍加改造：

def TextSearch(connection,text):
    text_conf = AutoConfig.load_config('text_image_embedding')
    text_conf.modality = 'text'

    text_pipe = AutoPipes.pipeline('text_image_embedding', text_conf)
    embedding = text_pipe(text).get()[0]
    
    res=SearchInTiDB(connection,embedding)
    p = (
        pipe.input('text','search_result')
        .map('search_result','prev',read_images)
        .output('text','prev')
    )
    DataCollection(p(text,res)).show()

還是用的clip_vit_base_patch16模型，只是使用模態改成了文字。透過對文字做 embedding 後得到向量資料送到 TiDB 中進行搜尋，流程和前面基本一樣。

看一下最終效果：

企業微信截圖_20240417104933.png

企業微信截圖_20240417105047.png

可以發現英文的搜尋效果要很多，這個主要是因為模型對於中文理解能力比較差，英文語義下 TiDB 的向量搜尋準確度依然非常高。

基於 TiDB Vector，前後不到100行程式碼就實現了以圖搜圖和自然語言搜圖。

未來展望

反正第一時間體驗完的感受就是：太香了，強烈推薦給大家！

在以往，想在關係型資料庫中對非結構化資料實現搜尋是一件不敢想象的事，哪怕是號稱無所不能的 PostgreSQL 在向量外掛的加持下也沒有獲得太多關注，這其中有場景、效能、生態等各方面的因素制約。而如今在 AI 大浪潮中，應用場景變得多樣化，生態鏈變得更豐富，TiDB Vector 的誕生恰逢其時。

但是不可忽視的是，傳統資料庫整合向量化的能力已經是大勢所趨，哪怕是 Redis 這樣的產品也擁有了向量能力。前有專門的向量資料庫阻擊，後有各種傳統資料庫追趕，這注定是一個慘烈的賽道，希望 TiDB 能深度打磨產品，突圍成功。

期待的功能：更多的索引型別、GPU加速等。

當然了，最大的願望必須是 TiDB On-Premises 中能儘快看到 Vector 的身影。

給 TiDB 點贊！

作者介紹：hey-hoho，來自神州數碼鈦合金戰隊，是一支致力於為企業提供分散式資料庫TiDB整體解決方案的專業技術團隊。團隊成員擁有豐富的資料庫從業背景，全部擁有TiDB高階資格證書，並活躍於TiDB開源社群，是官方認證合作伙伴。目前已為10+客戶提供了專業的TiDB交付服務，涵蓋金融、證券、物流、電力、政府、零售等重點行業。

相關文章