深度学习--FAISS向量数据库

简介

我们日常使用的各种 APP 中的许多功能，都离不开相似度检索技术。比如一个接一个的新闻和视频推荐、各种常见的对话机器人、保护我们日常账号安全的风控系统、能够用哼唱来找到歌曲的听歌识曲，甚至就连外卖配送的最佳路线选择也都有着它的身影。

相信很多同学是第一次听说它，或者只知道它的大名，而不知该如何使用它。本篇文章，我们就来聊聊 faiss，分享这个“黑科技”是如何发挥神奇的“魔法”的。

写在前面

faiss 是相似度检索方案中的佼佼者，是来自 Meta AI（原 Facebook Research）的开源项目，也是目前最流行的、效率比较高的相似度检索方案之一。虽然它和相似度检索这门技术颇受欢迎，在出现在了各种我们所熟知的“大厂”应用的功能中，但毕竟属于小众场景，有着不低的掌握门槛和复杂性。

所以，不要想着一口气就完全掌握它，咱们一步一步来。

当然，如果你实在懒得了解，希望能够和写简单的 Web 项目一样，写几行 CRUD 就能够完成高效的向量检索功能，可以试试启动一个 Milvus 实例。或者更懒一些的话，可以试着使用 Cloud 服务，来完成高性能的向量检索。

了解 Faiss 的工作机制和适用场景

在正式使用 faiss 之前，我们需要先了解它的工作机制。

当我们把通过模型或者 AI 应用处理好的数据喂给它之后（“一堆特征向量”），它会根据一些固定的套路，例如像传统数据库进行查询优化加速那样，为这些数据建立索引。避免我们进行数据查询的时候，需要笨拙的在海量数据中进行一一比对，这就是它能够实现“高性能向量检索”的秘密。

我们熟知的互联网企业中比较赚钱的“搜广推”（搜索、广告、推荐）业务中，会使用它解决这些场景下的向量召回工作。在这些场景下，系统需要根据多个维度进行数据关联计算，因为实际业务场景中数据量非常大，很容易形成类似“笛卡尔积”这种变态的结果，即使减少维度数量，进行循环遍历，来获取某几个向量的相似度计算，在海量数据的场景下也是不现实的。

而 Faiss 就是解决这类海量数据场景下，想要快速得到和查询内容相似结果（Top K 个相似结果），为数不多的靠谱方案之一。

和我们在常见数据库里指定字段类型一样， Faiss 也能够指定数据类型，比如 IndexFlatL2、IndexHNSW、IndexIVF等二十来种类型，虽然类型名称看起来比较怪，和传统的字符串、数字、日期等数据看起来不大一样，但这些场景将能够帮助我们在不同的数据规模、业务场景下，带来出乎意料的高性能数据检索能力。反过来说，在不同的业务场景、不同数据量级、不同索引类型和参数大小的情况下，我们的应用性能指标也会存在非常大的差异，如何选择合适的索引，也是一门学问。（下文会提到）

除了支持丰富的索引类型之外，faiss 还能够运行在 CPU 和 GPU 两种环境中，同时可以使用 C++ 或者 Python 进行调用，也有开发者做了 Go-Faiss ，来满足 Golang 场景下的 faiss 使用。

对 Faiss 有了初步认识之后，我们来进行 Faiss 使用的前置准备。

环境准备

为了尽可能减少不必要的问题，本篇文章中，我们使用 Linux 操作系统作为 faiss 的基础环境，同时使用 Python 作为和 faiss 交互的方式。

在之前的文章中，我介绍过如何准备 Linux 环境 和 Python 环境，如果你是 Linux 系统新手，可以阅读这篇文章，从零到一完成系统环境的准备：《在笔记本上搭建高性价比的 Linux 学习环境：基础篇》；如果你不熟悉 Python 的环境配置，建议阅读这篇文章《用让新海诚本人惊讶的 AI 模型制作属于你的动漫视频》，参考“准备工作”部分，完成 “Conda” 的安装配置。

在一切准备就绪之后，我们可以根据自己的设备状况，选择使用 CPU 版本的 faiss 还是 GPU 版本的 faiss，以及选择是否要指定搭配固定 CUDA 版本使用：

# 创建一个干净的环境
conda create -n faiss -y
# 激活这个环境
conda activate faiss
# 安装 CPU 版本
conda install -c pytorch python=3.8 faiss-cpu -y
# 或者，安装 GPU 版本
conda install -c pytorch python=3.8 faiss-gpu -y
# 或者，搭配指定 CUDA 版本使用
conda install -c pytorch python=3.8 faiss-gpu cudatoolkit=10.2 -y
1
2
3
4
5
6
7
8
9
10

在配置安装的时候，推荐使用 3.8 版本的 Python，避免不必要的兼容性问题。在准备好环境之后，我们就能够正式进入神奇的向量数据世界啦。

构建向量数据

前文提到了，适合 faiss 施展拳脚的地方是向量数据的世界，所以，需要先进行向量数据的构建准备。

本文作为入门篇，就先不聊如何对声音（音频）、电影（视频）、指纹和人脸（图片）等数据进行向量数据构建啦。我们从最简单的文本数据上手，实现一个“基于向量检索技术的文本搜索功能”。接下来，我将以我比较喜欢的小说 “哈利波特”为例，你可以根据自己的喜好调整要使用的文本数据。从网络上下载好要处理为向量的文本数据（txt 文档）。

简单针对数据进行 ETL

我这里的原始 TXT 文档尺寸是 3 MB 大小，为了减少不必要的向量转化计算量，我们先对内容进行必要的预处理（数据的 ETL 过程），去掉不必要的重复内容，空行等：

cat /Users/soulteary/《哈利波特》.txt | tr -d ' ' | sed '/^[[:space:]]*$/d' > data.txt
1

打开文本仔细观察，数据中有一些行中的文本数据格外长，是由好多个句子组成的，会对我们的向量特征计算、以及精准定位检索结果造成影响的。所以，我们还需要进行进一步的内容调整，将多个长句拆成每行一个的短句子。

为了更好的解决句子换行的问题，以及避免将一段人物对话中的多个句子拆散到多行，我们可以使用一段简单的 Node.js 脚本来处理数据：

const { readFileSync, writeFileSync } = require("fs");

const raw = readFileSync("./hp.txt", "utf-8")
  .split("\n")
  .map((line) => line.replace(/。/g, "。\n").split("\n"))
  .flat()
  .join("\n")
  .replace(/“([\S]+?)”/g, (match) => match.replace(/\n/g, ""))
  .replace(/“([\S\r\n]+?)”/g, (match) => match.replace(/[\r\n]/g, ""))
  .split("\n")
  .map((line) => line.replace(/s/g, "").trim().replace(/s/g, "—"))
  .filter((line) => line)
  .join("\n");

writeFileSync("./ready.txt", raw);
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

我们执行 node . 将文本处理完毕之后，当前文件夹中将出现一个名为 ready.txt 的文本文件。

为了方便后文中，我们更具象的了解向量数据库的资源占用，我们顺手查看下整理好的文本文件占磁盘空间是多少：

du -hs ready.txt 
5.5M	ready.txt
1
2

使用模型将文本转换为向量

为了将文本转换为向量数据，我们需要使用能够处理文本嵌入的模型。我这里选择的模型是来自人大、腾讯 AI Lab、北大（按论文作者顺序）联合推出的《UER: An Open-Source Toolkit for Pre-training Models》预训练模型。

关于这个预训练模型的相关资料：

• HuggingFace https://huggingface.co/uer/sbert-base-chinese-nli
• 训练数据 https://github.com/liuhuanyong/ChineseTextualInference/

想要使用模型，我们需要先安装一些 Python 的基础软件包：

pip install sentence_transformers pandas
1

在依赖安装完毕之后，我们可以在终端中输入 python 来进入 Python 交互式终端，首先将我们准备好的文本文件使用 pandas 解析为 DataFrames 。

import pandas as pd
df = pd.read_csv("ready.txt", sep="#",header=None, names=["sentence"])
print(df)
1
2
3

在执行之后，我们将看到类似下面的结果：

                                      sentence
0                                  《哈利波特》J.K罗琳
1                                第一部 第一章 幸存的男孩
2          住在四号普里怀特街的杜斯利先生及夫人非常骄傲地宣称自己是十分正常的人。
3      但是他们最不希望见到的就是任何奇怪或神秘故事中的人物因为他们对此总是嗤之以鼻。
4                      杜斯利先生是一家叫作格朗宁斯的钻机工厂的老板。
...                                        ...
60023                 哈利看着她茫然地低下头摸了摸额头上闪电形的伤疤。
60024                                “我知道他会的。”
60025                         十九年来哈利的伤疤再也没有疼过。
60026                                   一切都很好。
60027                                    （全书完）

[60028 rows x 1 columns]
1
2
3
4
5
6
7
8
9
10
11
12
13
14

接下来，我们对载入内存的文本进行向量计算，对每一行数据进行“特征向量抽取”：

from sentence_transformers import SentenceTransformer
model = SentenceTransformer('uer/sbert-base-chinese-nli')
sentences = df['sentence'].tolist()
sentence_embeddings = model.encode(sentences)
1
2
3
4

这个过程会比较久，消耗时间将会和你的电脑性能相关，我这边使用一台 Zen2 的普通笔记本，大概需要运行接近半个小时，所以这个时间不妨站起来动一动，缓解一天的疲劳。

当数据向量完毕之后，我们可以先执行 sentence_embeddings.shape，看看数据的状况：

(60028, 768)
1

执行完毕，我们将看到类似上面的结果，有六万条文本被向量化为了 768 维的向量数据。

在搞定向量数据之后，我们来实现第一个向量搜索的 Demo。

使用 faiss 实现最简单的向量检索功能

接下来，我们将使用 faiss 实现一个小功能，针对哈利波特小说全集内容，接触向量检索技术，完成相似内容搜索的功能。与我们使用“CTRL+F”或者把数据倒入 MySQL，使用 “%LIKE%” 去进行全文匹配不同，我们的工具性能，将会远远高于一般的检索方式。

为了能够得到“快到飞起”的执行效率，在使用 faiss 查询大量数据之前，我们首先需要和其他追求效率的数据库软件一样，为数据建立索引，我们先来看看最简单的平面索引：IndexFlatL2。

借助平面索引，完成基础的相似内容查询功能

faiss 中最简单的索引，便是没有使用任何花哨技巧（压缩、分区等）的平面索引：IndexFlatL2。当我们使用这种索引的时候，我们查询的数据会和索引中所有数据进行距离计算，获取它们之间的 L2 距离（欧几里得距离）。因为它会尽职尽责的和所有数据进行比对，所以它是所有索引类型中最慢的一种，但是也是最简单和最准确的索引类型，同时，因为类型简单，也是内存占用量最低的类型。而它采取的遍历式查找，也会被从业者打趣称之为“暴力搜索”。

在上文中，我们已经准备好了 768 维度的高维向量数据，接下来，我们就用这些数据来建立我们的第一堆“向量数据”的索引：

import faiss

dimension = sentence_embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(sentence_embeddings)
1
2
3
4
5

将我们数据的维度信息传递给 faiss.IndexFlatL2 函数，建立一个空的索引容器，然后使用 index.add(sentence_embeddings) 将我们在之前处理好的向量数据灌入这个索引容器中。

执行完毕上面的代码之后，我们执行 index.ntotal 来查看索引的数据是否正确：

# >>> index.ntotal
60028
1
2

确认所有数据都被索引之后，我们来写一段最简单的程序，来进行查询，为了演示“相似性检索”，而不是“关键词匹配”，我们来搜索一个离谱的原文肯定没有的内容“哈利波特猛然睡醒”：

topK = 5
search = model.encode(["哈利波特猛然睡醒"])
D, I = index.search(search, topK)
df['sentence'].iloc[I[0]]
1
2
3
4

执行程序之后，我们将能够看到比较符合预期的神奇结果：

# >>> topK = 5
# >>> search = model.encode(["哈利波特猛然睡醒"])
# >>> D, I = index.search(search, topK)
# >>> df['sentence'].iloc[I[0]]
38216  “我很好先生”哈利结巴的说擦着脸上的汗“我刚才只是睡着了做了个恶梦。
37890  “我很好先生”哈利结巴的说擦着脸上的汗“我刚才只是睡着了做了个恶梦。
8009   那天晚上哈利失眠了。
13996  最後哈利精疲力尽的爬上床猛拉他的四柱大床的蚊帐堵住射进来的一道月光翻身躺了进去而且几乎立即觉...
45306  罗恩立刻就进入了梦乡但是哈利在上床之前从行李箱里翻出了他的那本《高级魔药制备》。
Name: sentence, dtype: object
1
2
3
4
5
6
7
8
9
10

虽然没有完全匹配关键词，但是我们想要的内容还是被程序找到了。我们每天都在使用的搜索引擎背后的众多技术之一，也包括类似的向量检索（未来有机会的话，我们聊聊语义搜索）。

进一步了解向量检索的细节

我知道有一些同学，在惊叹上面这加起来不到 10 行的代码的效果之余，体验之后依旧对于“向量”的感知是零。为了让大家清清楚楚的“入坑”，我们先来展开聊聊上面的程序的含义。

topK = 5
search = model.encode(["哈利波特猛然睡醒"])
D, I = index.search(search, topK)
df['sentence'].iloc[I[0]]
1
2
3
4

第一行，topK 定义了我们要查找多少条最相似的数据，比如这里我就只想查询 5 条数据，避免有人说我水文章字数 本文内容由网友自发贡献，转载请注明出处：https://www.wpsshop.cn/w/代码探险家/article/detail/875012