TiDB Vector 抢先体验之用 TiDB 实现以图搜图

HOHO

发布于 2024-04-24 08:14:27

1400

发布于 2024-04-24 08:14:27

文章被收录于专栏：用户5447549的专栏

本文首发自 TiDB 社区专栏：https://tidb.net/blog/0c5672b9 转载请注明出处！

前言

最早知道 TiDB 要支持向量化的消息应该是在23年10月份左右，到第一次见到 TiDB Vector 的样子是在今年1月初，当时 dongxu 在朋友圈发了一张图：

去年我研究了一段时间的向量数据库，一直对 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 实例并没有太大区别，在创建集群页面可以看到加入了如下开关：

不过要注意的是目前 TiDB Vector 只在 AWS 的eu-central-1可用区开放，选到了其他可用区就看不到这个开关。

这里只需要填一个集群名称就可以开始创建，创建成功后的样子如下所示：

下面开始进入正题。

关于向量的那些事

一些基础概念

向量：向量就是一组浮点数，在编程语言中通常体现为 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维也是一样的方法。

以上内容来自我去年讲的向量数据库公开课：https://www.bilibili.com/video/BV1YP411t7Do

可以发现维数越多，对算力的要求就越高，计算时间就越长。

第一个 TiDB AI 应用：以图搜图

基础实现

借助前面介绍的理论知识，一个以图搜图的流程应该是这样子：

下面我用最简洁直白的代码演示整个流程，方便大家理解。

首先肯定是先连接到 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 框架的话这里对数据库和向量加工操作会简单些，不需要数组到字符串之间的手工转换。

加载完成后的数据：

下一步定义出根据指定向量在 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()

看一下最终搜索效果如何。先看一张已经在图片库存在的图（左边是待搜索的图，右边是搜索结果，按相似度由高到低）：

不能说非常相似，只能说是一模一样，准确度非常高！再看一下不在图片库的搜索效果：

图片库里有几十种动物，能够准确搜索出需要的是狗，特别是第一张从图片色彩、画面角度、动作神态上来说都非常相似。

使用向量索引优化

没错，向量也能加索引，但这个索引和传统的 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 中进行搜索，流程和前面基本一样。

看一下最终效果：