最近邻查找最优算法_最近邻搜索算法SPTAG的安装及使用

最新推荐文章于 2024-06-24 14:50:33 发布

weixin_39701861

最新推荐文章于 2024-06-24 14:50:33 发布

阅读量594

点赞数

文章标签：最近邻查找最优算法

本文链接：https://blog.csdn.net/weixin_39701861/article/details/112560457

版权

SPTAG是一个用于大规模矢量搜索的高效库，支持L2距离和余弦距离的比较。本文介绍了SPTAG的两种方法——SPTAG-KDT和SPTAG-BKT，并提供了Ubuntu 18.04上的Docker安装步骤。通过Dockerfile构建镜像，然后使用生成的镜像创建容器以运行SPTAG。此外，还概述了SPTAG的简单使用和接口程序，适用于近邻查找应用。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

SPTAG （Space Partition Tree And Graph）是分布式近似最近邻域搜索（ANN）库，为大规模矢量搜索场景提供高质量矢量索引构建，搜索和分布式在线服务工具包。

该库假定样本表示为向量，并且可以通过 L2 距离或余弦距离来比较向量。为查询向量返回的向量是与查询向量具有最小 L2 距离或余弦距离的向量。

SPTAG 提供两种方法：kd-tree 和相对邻域图（SPTAG-KDT），以及平衡 k-means 树和相对邻域图（SPTAG-BKT）。SPTAG-KDT 在指数构建成本方面是有利的，而 SPTAG-BKT 在非常高维数据中的搜索精度方面是有利的。

github地址：https://github.com/microsoft/SPTAG

安装步骤（环境Ubuntu18.04）：

推荐使用docker安装，不会的自行百度docker的基本使用。

首先将sptag库克隆到本地 , 执行git clone https://github.com/microsoft/SPTAG.git

查看目录文件，可以发现Dcokerfile,这就是用于生成sptag镜像的文件。

FROM ubuntu:18.04

WORKDIR /app
COPY CMakeLists.txt ./
COPY AnnService ./AnnService/
COPY Test ./Test/
COPY Wrappers ./Wrappers/

RUN apt-get update && apt-get -y install wget build-essential libtbb-dev 
    # remove the following if you don't want to build the wrappers
    openjdk-8-jdk python3-pip swig

# cmake >= 3.12 is required
RUN wget "https://github.com/Kitware/CMake/releases/download/v3.14.4/cmake-3.14.4-Linux-x86_64.tar.gz" -q -O - 
        | tar -xz --strip-components=1 -C /usr/local

# specific version of boost
RUN wget "https://dl.bintray.com/boostorg/release/1.67.0/source/boost_1_67_0.tar.gz" -q -O - 
        | tar -xz && 
        cd boost_1_67_0 && 
        ./bootstrap.sh && 
        ./b2 install && 
        # update ld cache so it finds boost in /usr/local/lib
        ldconfig && 
        cd .. && rm -rf boost_1_67_0

# build
RUN mkdir build && cd build && cmake .. && make && cd ..

# so python can find the SPTAG module
ENV PYTHONPATH=/app/Release

dockerfile的语法很简单，就是用于生成镜像的，这里只简单介绍下：

FROM ubuntu:18.04告诉你本镜像继承于ubuntu:18.04这个镜像，稍后会将怎么创建这个镜像。

WORKDIR /app 指定工作目录，相当于后续的操作都会在这个目录执行。

COPY就是复制的意思，将相关目录复制进工作目录。

RUN 执行相关命令，也就是/bin/bash的命令集，这里补充点， apt-get后面跟的这个-y是yes的意思，即在执行安装命令的时候遇到yes or no时，自动选择yes。另外如果你只需要python wrapper 则可以将openjdk-8-jdk注释掉，如果只需要java wrapper就将python3-pip swig注释掉。

ENV PYTHONPATH=/app/Release环境变量配置，这里主要告诉python编译器， sptag python封装库的位置，以便于可以加载python库。

这里不从下ubutnu：18.04镜像的生成。

首先从https://github.com/tianon/docker-brew-ubuntu-core/tree/3c462555392cb188830b7c91e29311b5fad90cfe/bionic这个地址去下载Dockerfile和ubuntu-bionic-core-cloudimg-amd64-root.tar.gz,将其放入同一文件夹，如/tmp,然后执行 docker build -t ubuntu:18.04 /tmp ,这里ubuntu:18.04是你创建的镜像名称， /tmp是Dockerfile文件存放的位置，执行完后你就可以通过sudo docker images看到你刚刚创建的镜像了。

这个时候再执行 sudo docker build -t sptag /github/SPTAG 即可完成sptag镜像的生成。

至此，就完成了sptag镜像的生成，接下来用它生成容器就可以使用了。

下面将简单说下sptag的简单使用和接口程序。

微软给的demo在test.py里，我自己主要用了metadata模式，并对其进行适当的修改，下面是我的主要接口函数：

import SPTAG
algo = 'BKT'
distmethod = 'L2'
def BuildWithMetaData(algo, distmethod, x, s, out):
    i = SPTAG.AnnIndex(algo, 'Float', x.shape[1])
    i.SetBuildParam("NumberOfThreads", '4')
    i.SetBuildParam("DistCalcMethod", distmethod)
    if i.BuildWithMetaData(x.tobytes(), s, x.shape[0]):
        i.Save(out)
        
def AddWithMetaData(textid, vector, index):
    i = SPTAG.AnnIndex.Load(index)
    k = 0
    while True:
        if i.AddWithMetaData(vector.tobytes(), textid, vector.shape[0]):
            i.Save(index)
            return True
        k += 1
        if k > 3:
            return False
        
def SearchWithMetaData(index, q, k):
    i = SPTAG.AnnIndex.Load(index)
    i.SetSearchParam("MaxCheck", '1024')
    result = i.SearchWithMetaData(q.tobytes(), k)
    out = {}
    for tt in range(k):
        out[result[2][tt].decode().replace("n", "")] = result[1][tt]
    return out

def DeleteData(index, x):
    i = SPTAG.AnnIndex.Load(index)
    k = 0
    while True:
        if i.Delete(x.tobytes(), x.shape[0]):
            i.Save(index)
            return True
        k += 1
        if k > 3:
            return False

这里的index实际就是数据库的目录，实际搜索中，在初始化函数调用i = SPTAG.AnnIndex.Load(index)，后面直接调用它的SearchWithMetaData方法就可以了。另外这个metadata, 我这是textid,是一个字符串，多条数据插入时，用'n'分割，python中还需调用s.encode()，将其转换为bytes类型，这里是c++中的一个bug，不知道现在修复没。

主要的安装及使用就是这些，至于如何将数据转换成向量，可以直接借助于现有的模型，如BERT，YOLO等，完成自己的应用。