跳转到主要内容

C++/Python中的近似最近邻搜索,优化内存使用和加载/保存到磁盘。

项目描述

https://img.shields.io/github/stars/spotify/annoy.svg

Annoy

Annoy example
https://github.com/spotify/annoy/actions/workflows/ci.yml/badge.svg

Annoy (近似最近邻 哇) 是一个C++库,具有Python绑定,用于搜索空间中与给定查询点接近的点。它还创建大型只读文件数据结构,这些数据结构被mmapped到内存中,以便许多进程可以共享相同的数据。

安装

要安装,只需执行pip install --user annoyPyPI下载最新版本。

C++版本只需克隆仓库并#include "annoylib.h"

背景

存在一些其他库可以进行最近邻搜索。Annoy 几乎与最快的库一样快(见下文),但它实际上还有一个真正让 Annoy 独具特色的功能:它可以 使用静态文件作为索引。特别是,这意味着你可以 在进程间共享索引。Annoy 还将创建索引与加载它们解耦,因此你可以将索引作为文件传递,并快速将其映射到内存中。Annoy 的另一个优点是它试图最小化内存占用,因此索引相当小。

这有什么用?如果你想查找最近邻,而你有很多 CPU,你只需要构建索引一次。你还可以在生产环境、Hadoop 作业等中传递和分发静态文件。任何进程都将能够将(mmap)索引加载到内存中,并立即进行查找。

我们在Spotify上使用它进行音乐推荐。在运行矩阵分解算法后,每个用户/项目都可以表示为 f 维空间中的一个向量。这个库帮助我们搜索相似的用户/项目。我们有一个高维空间中的数百万首曲目,因此内存占用是一个主要关注点。

Annoy 是由Erik BernhardssonHack Week期间的一两天内构建的。

功能概述

  • 欧几里得距离曼哈顿距离余弦距离汉明距离点积(内积)距离

  • 余弦距离等于归一化向量的欧几里得距离 = sqrt(2-2*cos(u, v))

  • 如果你的维度不多(如 <100),效果会更好,但似乎即使高达 1,000 维也能出奇地好

  • 内存占用小

  • 允许你在多个进程之间共享内存

  • 索引创建与查找分离(特别是一旦创建了树,就不能再添加更多项目)

  • 原生 Python 支持,已测试 2.7、3.6 和 3.7 版本。

  • 在磁盘上构建索引以支持索引无法装入内存的大数据集(由Rene Hollander贡献)

Python 代码示例

from annoy import AnnoyIndex
import random

f = 40  # Length of item vector that will be indexed

t = AnnoyIndex(f, 'angular')
for i in range(1000):
    v = [random.gauss(0, 1) for z in range(f)]
    t.add_item(i, v)

t.build(10) # 10 trees
t.save('test.ann')

# ...

u = AnnoyIndex(f, 'angular')
u.load('test.ann') # super fast, will just mmap the file
print(u.get_nns_by_item(0, 1000)) # will find the 1000 nearest neighbors

目前它只接受整数作为项目的标识符。请注意,它将为 max(id)+1 个项目分配内存,因为它假定你的项目编号为 0 … n-1。如果你需要其他标识符,你必须自己跟踪映射。

完整的 Python API

  • AnnoyIndex(f, metric) 返回一个新的可读写索引,存储 f 维的向量。度量可以是 "angular""euclidean""manhattan""hamming""dot"

  • a.add_item(i, v) 添加项目 i(任何非负整数)和向量 v。请注意,它将为 max(i)+1 个项目分配内存。

  • a.build(n_trees, n_jobs=-1) 构建 n_trees 棵树的森林。更多的树在查询时提供更高的精度。在调用 build 之后,不能再添加更多项目。 n_jobs 指定用于构建树的线程数。 n_jobs=-1 使用所有可用的 CPU 核心。

  • a.save(fn, prefault=False) 将索引保存到磁盘并加载它(请参阅下一个函数)。在保存后,不能再添加更多项目。

  • a.load(fn, prefault=False) 从磁盘加载(mmaps)索引。如果 prefault 设置为 True,则将整个文件预先读取到内存中(使用 mmap 与 MAP_POPULATE)。默认为 False

  • a.unload() 卸载。

  • a.get_nns_by_item(i, n, search_k=-1, include_distances=False) 返回最近的 n 个项目。在查询过程中,它将检查最多 search_k 个节点,如果未提供,则默认为 n_trees * nsearch_k 在运行时提供了更好的准确性和速度之间的权衡。如果将 include_distances 设置为 True,则将返回一个包含两个列表的 2 元素元组:第二个列表包含所有对应的距离。

  • a.get_nns_by_vector(v, n, search_k=-1, include_distances=False) 与此相同,但通过向量 v 进行查询。

  • a.get_item_vector(i) 返回之前添加的项目 i 的向量。

  • a.get_distance(i, j) 返回项目 ij 之间的距离。注意:这曾经返回的是平方距离,但从 2016 年 8 月起已更改。

  • a.get_n_items() 返回索引中的项目数量。

  • a.get_n_trees() 返回索引中的树的数量。

  • a.on_disk_build(fn) 准备将 Annoy 的索引构建在指定的文件中而不是 RAM 中(在添加项目之前执行,无需在构建后保存)

  • a.set_seed(seed) 将使用给定的种子初始化随机数生成器。仅用于构建树,即只在添加项目之前传递此参数。在调用 a.build(n_trees)a.load(fn) 之后将没有效果。

注意

  • 没有对值执行边界检查,因此请小心。

  • Annoy 使用归一化向量的欧几里得距离作为其角距离,对于两个向量 u,v 等于 sqrt(2(1-cos(u,v)))

C++ API 非常相似:只需 #include "annoylib.h" 即可访问它。

权衡

Annoy 只需要两个主要参数进行调整:树的数量 n_trees 和搜索时检查的节点数量 search_k

  • n_trees 在构建时提供,会影响构建时间和索引大小。较大的值将提供更准确的结果,但索引更大。

  • search_k 在运行时提供,会影响搜索性能。较大的值将提供更准确的结果,但将需要更长的时间来返回。

如果未提供 search_k,则默认为 n * n_trees,其中 n 是近似最近邻的数量。否则,search_kn_trees 大致独立,即如果保持 search_k 不变,则 n_trees 的值不会影响搜索时间,反之亦然。基本上,建议根据您能够承受的内存量将 n_trees 设置得尽可能大,并且建议根据您对查询的时间限制将 search_k 设置得尽可能大。

您也可以为了减少加载时间、内存使用和磁盘I/O,接受较慢的搜索时间。在支持的平台上,在加载和保存期间,索引是预先缓存的,导致文件预先从磁盘读取到内存中。如果您将 prefault 设置为 False,则 mmapped 索引的页面将根据需要从磁盘读取并在内存中缓存,以完成搜索。这可能会显著增加早期搜索时间,但对于内存小于索引大小的系统,当对已加载索引执行少量查询时,以及/或者当索引的大面积不太可能对搜索查询相关时,可能更适合。

它如何工作

通过使用 随机投影 并构建树。在树的每个中间节点,选择一个随机超平面,将空间分为两个子空间。这个超平面是通过从子集中采样两个点并取等距于它们的超平面来选择的。

我们重复此过程k次,以得到一个树丛。k需要根据您的需求进行调整,通过查看您在精度和性能之间的权衡。

汉明距离(由 Martin Aumüller 贡献)在底层将数据打包到64位整数中,并使用内置的位计数原语,因此它可能相当快。所有分割都是轴对齐的。

点积距离(由 Peter Sobot 贡献)使用 微软研究院 Bachrach 等人于 2014 年发表的方法,将提供的向量从点(或“内积”)空间降低到更查询友好的余弦空间。

更多信息

ANN benchmarks

源代码

所有内容都是用 C++ 编写的,附带一些针对性能和内存使用的优化。已经警告过了 :)

代码应支持Windows,归功于强口蒂莫西·瑞利

要运行测试,请执行python setup.py nosetests。测试套件包括一个从互联网下载的大规模真实世界数据集,因此执行将需要几分钟。

讨论

欢迎将任何问题或评论发布到annoy-user群组。我在Twitter上的昵称是@fulhack

项目详情


下载文件

下载适合您平台的应用程序。如果您不确定选择哪个,请了解更多关于安装包的信息。

源分发

annoy-1.17.3.tar.gz (647.5 kB 查看哈希值)

上传时间

构建分发

annoy-1.17.3-cp310-cp310-macosx_11_0_arm64.whl (57.7 kB 查看哈希值)

上传时间 CPython 3.10 macOS 11.0+ ARM64

支持者