基于Qdrant和FiftyOne的最近邻嵌入搜索算法实战

译者 | 朱先忠

审校 | 孙淑娟

简介

神经网络嵌入能够实现输入数据的低维表示，从而服务于各种类型的神经网络应用程序。嵌入有一些有趣的功能，因为它们可以捕获数据点的语义。这对于图像和视频等非结构化数据尤其有用，这样不仅可以对像素相似性进行编码，还可以对一些更复杂的关系进行编码。

使用FiftyOne和Plotly可视化针对BDD100K数据集的嵌入

在这些嵌入上执行搜索可以应用于许多场景下，如分类任务、建立推荐系统，甚至是异常检测任务。对嵌入执行最近邻搜索以完成这些任务的主要好处之一是，无需为每个新问题创建自定义网络；您可以经常使用预先训练的模型。而且，使用某些公开可用模型生成的嵌入是完全可能的，而无需任何进一步的微调。

虽然有很多涉及嵌入的强大使用场景，但在执行嵌入搜索的工作流中往往都存在一些不同程度的挑战。具体而言，在大型数据集上执行最近邻搜索，然后有效地对搜索结果采取行动——例如执行自动标记数据等工作流，都存在技术和工具方面的挑战。为此，借助于Qdrant和FiftyOne这两个开源软件可以帮助简化这些工作流任务。

• Qdrant是一个开源的向量数据库，旨在对密集的神经嵌入执行近似最近邻搜索（ANN），这对于任何预计将扩展到大数据量的生产就绪系统都是非常必要的。
• FiftyOne是一个开源数据集管理和模型评估工具，允许开发人员有效地管理和可视化数据集，生成嵌入并改进模型结果。

在本文中，我们将MNIST数据集加载到FiftyOne中，并基于ANN进行分类。数据点将通过从训练数据集中的K个最近点中选择最常见的参考答案标签进行分类。换句话说，对于每个测试样本，我们将使用选定的距离函数选择其K个最近的邻居，然后通过投票选择最佳标签。另一方面，向量空间中的所有搜索都将使用Qdrant来完成，以加快速度。然后，我们将在FiftyOne中评估该分类的结果。

安装

如果您想开始使用Qdrant的语义搜索，您需要运行它的一个实例，因为该工具以客户机—服务器方式工作。要做到这一点，最简单的方法是使用一个官方Docker映像，只需一个命令即可启动Qdrant：

docker run -p“6333:6333”-p“63 34:6334”-d qdrant/qdrant

运行该命令后，我们将运行Qdrant服务器，端口6333处暴露HTTP API，端口6334处暴露gRPC接口。

此外，我们还需要安装一些Python包。我们将使用FiftyOne来可视化数据、它们的参考答案标签以及嵌入相似性模型预测的数据。嵌入将通过torchvision中提供的MobileNet v2创建。当然，我们也需要以某种方式与Qdrant服务器通信，因为我们将使用Python，所以Qdrant_client是实现这一任务的首选方式。

pip install fiftyone
pip install torchvision
pip install qdrant_client

整体任务流程

1. 加载数据集

2. 生成嵌入

3. 将嵌入加载到Qdrant

4. 最近邻分类

5. FiftyOne评估

加载数据集

为了使事情顺利进行，我们需要采取几个步骤。首先，我们需要加载​​MNIST数据集​​，并从中提取训练样本，因为我们将在搜索操作中使用它们。为了使一切更快，我们不打算使用所有的样本，只使用2500个样本。我们可以使用​​FiftyOne数据集Zoo​​，并通过一行代码来加载所需的MNIST子集。

import fiftyone as fo
import fiftyone.zoo as foz

# 加载数据
dataset = foz.load_zoo_dataset("mnist", max_samples=2500)

#取得所有的训练样本
train_view = dataset.match_tags(tags=["train"])

接下来，让我们仔细分析一下FiftyOne应用程序中所使用的数据集。

# 可视化FiftyOne中的数据集
session = fo.launch_app(train_view)

生成嵌入

下一步是在数据集中的样本上生成嵌入。这始终可以在FiftyOne之外使用您的自定义模型完成。然而，FiftyOne也在其中的FiftyOneModelZoo中提供了各种模型，这些模型可以直接用于生成嵌入。

在本例中，我们使用在ImageNet上训练的MobileNetv2来计算每个图像的嵌入。

#计算每个图像的嵌入
model = foz.load_zoo_model("mobilenet-v2-imagenet-torch")
train_embeddings = train_view.compute_embeddings(model)

将嵌入结果加载到Qdrant

Qdrant不仅可以存储向量，还可以存储一些相应的属性——每个数据点都有一个相关的向量，还可以选择附带一个JSON类型的属性项。我们想用这个来传递参考答案标签，以确保我们可以在稍后做出预测。

ground_truth_labels = train_view.values("ground_truth.label")
train_payload = [
    {"ground_truth": gt} for gt in ground_truth_labels
]

创建嵌入后，我们可以开始与Qdrant服务器进行通信了。需要说明的是，QdrantClient有一个类实例比较有用，因为它包含了所有必需的方法。让我们连接并创建一个名为“mnist”的点集合。向量大小取决于模型输出；因此，如果我们想在以后使用不同的模型进行实验，我们需要导入不同的模型，但其余的将保持不变。最终，在确保集合存在后，我们可以发送所有向量及其包含其真实标签的有效载荷。

import qdrant_client as qc
from qdrant_client.http.models import Distance

# 把已经训练的嵌入加载到Qdrant
def create_and_upload_collection(
    embeddings, payload, collection_name="mnist"
):
    client = qc.QdrantClient(host="localhost")
    client.recreate_collection(
        collection_name=collection_name,
        vector_size=embeddings.shape[1],
        distance=Distance.COSINE,
    )
    client.upload_collection(
        collection_name=collection_name,
        vectors=embeddings,
        payload=payload,
    )
    return client
    
client = create_and_upload_collection(train_embeddings, train_payload)

最近邻分类

现在对数据集执行推断。我们可以为测试数据集创建嵌入，但忽略基本事实，并尝试使用ANN找到它，然后比较两者是否匹配。让我们一步一步地从创建嵌入开始。

#通过在每个样本的邻居中选择最常见的标签，将标签分配给测试嵌入
test_view = dataset.match_tags(tags=["test"])
test_embeddings = test_view.compute_embeddings(model)

是时候“施展魔法”了。让我们遍历测试数据集的样本和相应的嵌入，并使用搜索操作从训练集中找到15个最接近的嵌入。我们还需要选择有效载荷，因为它们包含找到特定点附近最常见标签所需的参考答案标签。借助于Python的Counter类，我们可以避免再重写任何样板代码。最常见的标签将作为“ann_production”存储在FiftyOne中的每个测试样本上。

这些内容都体现在下面这个函数中。该函数将使用嵌入向量作为输入，并使用Qdrant搜索功能查找与测试嵌入最近的邻居，生成类型预测，并返回一个FiftyOne Classification对象，我们可以将其存储在FiftyOne数据集中。

import collections
from tqdm import tqdm

def generate_fiftyone_classification(
    embedding, collection_name="mnist"
):
    search_results = client.search(
        collection_name=collection_name,
        query_vector=embedding,
        with_payload=True,
        top=15,
    )
    # 统计每个类型的出现次数，并选择最常见的标签，置信度估计为最常见标签的出现次数除以结果总数
    counter = collections.Counter(
        [point.payload["ground_truth"] for point in search_results]
    )
    predicted_class, occurences_num = counter.most_common(1)[0]
    confidence = occurences_num / sum(counter.values())
    prediction = fo.Classification(
        label=predicted_class, confidence=confidence
    )
    return prediction
    
predictions = []

#调用Qdrant查找最近的数据点
for embedding in tqdm(test_embeddings):
    prediction = generate_fiftyone_classification(embedding)
    predictions.append(prediction)
    
test_view.set_values("ann_prediction", predictions)

我们通过计算属于最常见标签的样本分数来估计置信度。这给了我们一种直觉，即我们在预测每种情况的标签时有多么确定，并且可以在FiftyOne中使用，轻松找出令人困惑的样本。

FiftyOne评估

现在是取得一些成果的时候了！让我们从可视化这个分类器的表现开始。我们可以轻松启动​​FiftyOne应用程序​​来查看参考答案标签、预测结果和图像。

session = fo.launch_app(test_view)

FiftyOne提供了各种内置方法来评估模型预测，包括图像和视频数据集上的回归、分类、检测、多边形、实例和语义分割。通过下面两行代码，我们即可以计算并打印分类器的评估报告。

# 根据ground_truth中的值评估ANN预测
results = test_view.evaluate_classifications(
    "ann_prediction", gt_field="ground_truth", eval_key="eval_simple"
)

#显示分类指标
results.print_report()

在评估FiftyOne之后，我们可以使用results对象生成一个交互式混淆矩阵（https://voxel51.com/docs/fiftyone/user_guide/plots.html#confusion-matrices），允许我们点击单元格并自动更新应用程序以显示相应的样本。

plot = results.plot_confusion_matrix()
plot.show()

我们还可以再深入一点。我们可以使用FiftyOne复杂的​​查询语言​​，轻松找到所有与实际情况不匹配的预测，但预测的可信度很高。这些通常是数据集中最令人困惑的样本，不过也是我们可以从中获得最深刻见解的样本。

from fiftyone import ViewField as F

# 显示FiftyOne应用程序，但仅包括有相当把握预测出的错误预测
false_view = (
    test_view
    .match(F("eval_simple") == False)
    .filter_labels("ann_prediction", F("confidence") > 0.7)
)
session.view = false_view

上图中展示了模型中最令人困惑的样本，正如您所看到的，与数据集中的其他图像相比，它们非常不规则。我们可以采取的改进模型性能的下一步可能是使用FiftyOne来添加与这些类似的精确样本。然后，可以通过FiftyOne与CVAT，以及Labelbox等工具之间的集成对这些样本进行注释。此外，我们可以使用更多的向量进行训练，或者通过相似性学习对模型进行微调，例如，使用三元组损失算法。但现在，这个使用FiftyOne和Quadrant进行向量相似性分类的示例已经能够很好地工作了。

就这么简单，我们使用FiftyOne和Qdrant作为嵌入后端创建了一个ANN分类模型，因此查找向量之间的相似性可以不再像传统k-NN那样成为我们任务的瓶颈了。

请自己试试看

最后，Github代码仓库中的​​笔记本文件​​包含了你在本文中看到的所有内容的源代码。此外，它还包括该过程的一个实际用例，以在BDD100K道路场景数据集上执行夜间和白天属性的预注释。

总之，FiftyOne和Qdrant两个开源库可以一起使用，有效地对嵌入执行最近邻居搜索，并对图像和视频数据集的结果进行操作。这个过程的美妙之处在于它的灵活性和可重复性。您可以轻松地将新字段的附加参考答案标签加载到FiftyOne和Qdrant中，并使用现有嵌入重复此预注释过程。这可以快速降低注释成本，并更快地生成更高质量的数据集。

译者介绍

朱先忠，51CTO社区编辑，51CTO专家博客、讲师，潍坊一所高校计算机教师，自由编程界老兵一枚。

原文标题：​​Nearest Neighbor Embeddings Search With Qdrant and FiftyOne​​，作者：Eric Hofesmann