nmslib/hnswlib

GitHub: nmslib/hnswlib

一个轻量级的 header-only C++ 库,提供快速近似最近邻搜索功能,并附带 Python 绑定,支持增量索引构建、元素增删更新与持久化。

Stars: 5281 | Forks: 833

# Hnswlib - 快速近似最近邻搜索 带有 Python 绑定、插入和更新功能的 Header-only C++ HNSW 实现。 **新闻:** **版本 0.9.0** * 修复了带过滤器的暴力搜索中结果不正确的问题 (#514),由 [@lukaszsmolinski](https://github.com/lukaszsmolinski) 贡献 * 修复了 BFIndex 中缺失归一化检查的问题 (#514),由 [@lukaszsmolinski](https://github.com/lukaszsmolinski) 贡献 * 在可用元素少于 k 个时抛出异常 (#514),由 [@lukaszsmolinski](https://github.com/lukaszsmolinski) 贡献 * 移除了未使用的变量 (#531),由 [@lulyon](https://github.com/lulyon) 贡献 * 在 README 中将余弦相似度改为余弦距离,由 [@yurymalkov](https://github.com/yurymalkov) 贡献 **版本 0.8.0** * 多向量文档搜索和 epsilon 搜索(目前仅支持 C++) * 默认情况下不进行统计聚合,从而加快了多线程搜索的速度(看起来似乎也没有人使用它:[Issue #495](https://github.com/nmslib/hnswlib/issues/495))。 * 各种 bug 修复和改进 * `get_items` 现在具有 `return_type` 参数,其值可以是 'numpy' 或 'list' 完整的变更列表:https://github.com/nmslib/hnswlib/pull/523 **版本 0.7.0** * 添加了过滤支持 (#402, #430),由 [@kishorenc](https://github.com/kishorenc) 贡献 * 添加了用于过滤的 Python 接口(但请注意其性能受限于 GIL)(#417),由 [@gtsoukas](https://github.com/gtsoukas) 贡献 * 添加了用新插入的元素替换被标记为删除的元素的支持(用于控制索引大小,#418),由 [@dyashuni](https://github.com/dyashuni) 贡献 * 修复了更新/插入中的数据竞争/死锁问题,添加了针对多线程操作的压测 (#418),由 [@dyashuni](https://github.com/dyashuni) 贡献 * 文档、测试、异常处理、重构 (#375, #379, #380, #395, #396, #401, #406, #404, #409, #410, #416, #415, #431, #432, #433),由 [@jlmelville](https://github.com/jlmelville)、[@dyashuni](https://github.com/dyashuni)、[@kishorenc](https://github.com/kishorenc)、[@korzhenevski](https://github.com/korzhenevski)、[@yoshoku](https://github.com/yoshoku)、[@jianshu93](https://github.com/jianshu93)、[@PLNech](https://github.com/PLNech) 贡献 * 全局链接 (#383),由 [@MasterAler](https://github.com/MasterAler) 贡献,以及在 MSVC 中使用 USE_SSE (#408),由 [@alxvth](https://github.com/alxvth) 贡献 ### 核心亮点: 1. 轻量级,header-only,除 C++ 11 外无其他依赖 2. 提供 C++、Python 接口,外部支持 Java 和 R (https://github.com/jlmelville/rcpphnsw)。 3. 全面支持增量索引构建和元素更新(感谢 Apoorv Sharma 的贡献)。支持元素删除 (通过在索引中标记它们,之后可以被其他元素替换)。Python 索引支持 pickle 序列化。 4. 可以使用自定义的用户定义距离(C++)。 5. 与当前的 nmslib 实现相比,内存占用显著减少且构建时间更快。 有关算法参数的描述可以在 [ALGO_PARAMS.md](ALGO_PARAMS.md) 中找到。 ### Python 绑定 #### 支持的距离: | 距离 | 参数 | 公式 | | ------------- |:---------------:| -----------------------:| |平方 L2 (Squared L2) |'l2' | d = sum((Ai-Bi)^2) | |内积 (Inner product) |'ip' | d = 1.0 - sum(Ai\*Bi) | |余弦距离 (Cosine distance) |'cosine' | d = 1.0 - sum(Ai\*Bi) / sqrt(sum(Ai\*Ai) * sum(Bi\*Bi))| 请注意,内积并不是真正的度量标准。一个元素可能与另一个元素的距离比与自身的距离还要近。如果您将所有与自身不是最接近的元素从索引中移除,这可以带来一定的速度提升。 对于其他空间,请使用 nmslib 库 https://github.com/nmslib/nmslib。 #### API 描述 * `hnswlib.Index(space, dim)` 在空间 `space` 中创建一个维度为整数 `dim` 的未初始化 HNSW 索引。 `hnswlib.Index` 方法: * `init_index(max_elements, M = 16, ef_construction = 200, random_seed = 100, allow_replace_deleted = False)` 初始化一个没有任何元素的索引。 * `max_elements` 定义了结构中可存储元素的最大数量(可以增加/缩减)。 * `ef_construction` 定义了构建时间与准确率之间的权衡(参见 [ALGO_PARAMS.md](ALGO_PARAMS.md))。 * `M` 定义了图中传出连接的最大数量([ALGO_PARAMS.md](ALGO_PARAMS.md))。 * `allow_replace_deleted` 允许用新增的元素替换被删除的元素。 * `add_items(data, ids, num_threads = -1, replace_deleted = False)` - 将 `data`(形状为 `N*dim` 的 numpy 向量数组)插入到结构中。 * `num_threads` 设置要使用的 CPU 线程数(-1 表示使用默认值)。 * `ids` 是可选的大小为 N 的 numpy 数组,表示 `data` 中所有元素的整数标签。 - 如果索引中已经包含具有相同标签的元素,它们的特征将被更新。请注意,更新过程比插入新元素慢,但在内存和查询效率上更高。 * `replace_deleted` 替换被删除的元素。请注意,这样可以节省内存。 - 要使用此功能,调用 `init_index` 时必须传入 `allow_replace_deleted=True` * 与其他 `add_items` 调用是线程安全的,但与 `knn_query` 不是。 * `mark_deleted(label)` - 将元素标记为已删除,因此它将在搜索结果中被省略。如果该元素已被删除,则抛出异常。 * `unmark_deleted(label)` - 取消元素的删除标记,因此它将不会被排除在搜索结果之外。 * `resize_index(new_size)` - 更改索引的最大容量。与 `add_items` 和 `knn_query` 不是线程安全的。 * `set_ef(ef)` - 设置查询时间准确率与速度的权衡,由 `ef` 参数定义( [ALGO_PARAMS.md](ALGO_PARAMS.md))。请注意,该参数目前不会随索引一起保存,因此您需要在加载后手动设置。 * `knn_query(data, k = 1, num_threads = -1, filter = None)` 为每个元素的 `data`(形状为 `N*dim`)进行批量查询,查找 `k` 个最接近的元素。 * 返回一个形状为(`N*k`)的 numpy 数组。 * `num_threads` 设置要使用的 CPU 线程数(-1 表示使用默认值)。 * `filter` 根据标签过滤元素,返回具有允许 ID 的元素。请注意,在 Python 的多线程模式下,带过滤器的搜索运行速度较慢。建议设置 `num_threads=1` * 与其他 `knn_query` 调用是线程安全的,但与 `add_items` 不是。 * `load_index(path_to_index, max_elements = 0, allow_replace_deleted = False)` 将持久化的索引加载到未初始化的索引中。 * `max_elements`(可选)重置结构中元素的最大数量。 * `allow_replace_deleted` 指定正在加载的索引是否启用了替换被删除元素的功能。 * `save_index(path_to_index)` 将索引保存到持久化存储中。 * `set_num_threads(num_threads)` 设置数据插入/查询期间使用的默认 CPU 线程数。 * `get_items(ids, return_type = 'numpy')` - 如果 `return_type` 为 `list`,则返回一个形状为 `N*dim` 的向量列表;否则返回一个形状为(`N*dim`)的 numpy 数组,这些向量具有 `ids` numpy 向量(形状为 `N`)中指定的整数标识符。请注意,对于余弦相似度,它目前返回的是**归一化**后的向量。 * `get_ids_list()` - 返回所有元素 ID 的列表。 * `get_max_elements()` - 返回索引的当前容量 * `get_current_count()` - 返回当前索引中存储的元素数量 `hnswlib.Index` 类的只读属性: * `space` - 空间的名称(可以是 "l2"、"ip" 或 "cosine" 之一)。 * `dim` - 空间的维度。 * `M` - 定义图中传出连接最大数量的参数。 * `ef_construction` - 控制索引构建期间速度/准确率权衡的参数。 * `max_elements` - 索引的当前容量。等同于 `p.get_max_elements()`。 * `element_count` - 索引中的项目数量。等同于 `p.get_current_count()`。 `hnswlib.Index` 中支持读写的属性: * `ef` - 控制查询时间与准确率权衡的参数。 * `num_threads` - 在 `add_items` 或 `knn_query` 中使用的默认线程数。请注意,调用 `p.set_num_threads(3)` 等同于 `p.num_threads=3`。 #### Python 绑定示例 [在此处查看更多示例](examples/python/EXAMPLES.md): * 创建索引、插入元素、搜索、序列化/反序列化 * 在搜索期间使用布尔函数进行过滤 * 删除元素并将已删除元素的内存重新用于新添加的元素 创建索引、插入元素、搜索和 pickle 序列化的示例: ``` import hnswlib import numpy as np import pickle dim = 128 num_elements = 10000 # 生成样本数据 data = np.float32(np.random.random((num_elements, dim))) ids = np.arange(num_elements) # 声明 index p = hnswlib.Index(space = 'l2', dim = dim) # possible options are l2, cosine or ip # 初始化 index - 应事先知道元素的最大数量 p.init_index(max_elements = num_elements, ef_construction = 200, M = 16) # 元素插入(可以多次调用): p.add_items(data, ids) # 通过设置 ef 来控制召回率: p.set_ef(50) # ef should always be > k # 查询数据集,k - 最近元素的数量(返回 2 个 numpy 数组) labels, distances = p.knn_query(data, k = 1) # Index 对象支持 pickling # 警告:通过 pickle.dumps(p) 或 p.__getstate__() 的序列化与 p.add_items 方法一起使用时不是线程安全的! # 注意:ef 参数包含在序列化中;在 Index 加载时使用 random_seed 初始化随机数生成器 p_copy = pickle.loads(pickle.dumps(p)) # creates a copy of index p using pickle round-trip ### Index 参数作为类属性公开: print(f"Parameters passed to constructor: space={p_copy.space}, dim={p_copy.dim}") print(f"Index construction: M={p_copy.M}, ef_construction={p_copy.ef_construction}") print(f"Index size is {p_copy.element_count} and index capacity is {p_copy.max_elements}") print(f"Search speed/quality trade-off parameter: ef={p_copy.ef}") ``` 序列化/反序列化后进行更新的示例: ``` import hnswlib import numpy as np dim = 16 num_elements = 10000 # 生成样本数据 data = np.float32(np.random.random((num_elements, dim))) # 我们将数据分成两批: data1 = data[:num_elements // 2] data2 = data[num_elements // 2:] # 声明 index p = hnswlib.Index(space='l2', dim=dim) # possible options are l2, cosine or ip # 初始化 index # max_elements - 元素的最大数量(容量)。如果超出 # 在插入元素时将抛出异常。 # 可以通过保存/加载 index 来增加容量,见下文。 # # ef_construction - 控制 index 搜索速度/构建速度的权衡 # # M - 与数据的内部维度紧密相关。强烈影响内存消耗(~M) # 在固定的 ef/efConstruction 下,更高的 M 会带来更高的准确性/运行时间 p.init_index(max_elements=num_elements//2, ef_construction=100, M=16) # 通过设置 ef 来控制召回率: # 更高的 ef 会带来更好的准确性,但搜索速度更慢 p.set_ef(10) # 设置批量搜索/构建期间使用的线程数 # 默认使用所有可用的核心 p.set_num_threads(4) print("Adding first batch of %d elements" % (len(data1))) p.add_items(data1) # 查询元素自身并测量召回率: labels, distances = p.knn_query(data1, k=1) print("Recall for the first batch:", np.mean(labels.reshape(-1) == np.arange(len(data1))), "\n") # 序列化和删除 index: index_path='first_half.bin' print("Saving index to '%s'" % index_path) p.save_index("first_half.bin") del p # 重新初始化,加载 index p = hnswlib.Index(space='l2', dim=dim) # the space can be changed - keeps the data, alters the distance function. print("\nLoading index from 'first_half.bin'\n") # 增加总容量(max_elements),以便它能处理新数据 p.load_index("first_half.bin", max_elements = num_elements) print("Adding the second batch of %d elements" % (len(data2))) p.add_items(data2) # 查询元素自身并测量召回率: labels, distances = p.knn_query(data, k=1) print("Recall for two batches:", np.mean(labels.reshape(-1) == np.arange(len(data))), "\n") ``` #### C++ 示例 [在此处查看示例](examples/cpp/EXAMPLES.md): * 创建索引、插入元素、搜索、序列化/反序列化 * 在搜索期间使用布尔函数进行过滤 * 删除元素并将已删除元素的内存重新用于新添加的元素 * 多线程使用 * 多向量搜索 * Epsilon 搜索 ### 绑定安装 您可以从源码安装: ``` apt-get install -y python-setuptools python-pip git clone https://github.com/nmslib/hnswlib.git cd hnswlib pip install . ``` 或者您可以通过 pip 安装: `pip install hnswlib` ### 面向开发者 进行更改时,请运行测试(如果有新功能,请在 `tests/python` 中添加测试): ``` python -m unittest discover --start-directory tests/python --pattern "bindings_test*.py" ``` ### 其他实现 * 非度量空间库 (nmslib) - 主库(Python, C++),支持特殊的距离度量:https://github.com/nmslib/nmslib * 由 Facebook 开发的 Faiss 库,使用自己的 HNSW 实现进行粗量化(Python, C++): https://github.com/facebookresearch/faiss * 以下论文的代码 ["Revisiting the Inverted Indices for Billion-Scale Approximate Nearest Neighbors"](https://arxiv.org/abs/1802.02422) (目前在压缩索引领域处于最先进水平,C++): https://github.com/dbaranchuk/ivf-hnsw * Amazon PECOS https://github.com/amzn/pecos * TOROS N2 (Python, C++):https://github.com/kakao/n2 * Online HNSW (C++):https://github.com/andrusha97/online-hnsw) * Go 实现:https://github.com/Bithack/go-hnsw * Python 实现(作为由 Matteo Dell'Amico 编写的聚类代码的一部分):https://github.com/matteodellamico/flexible-clustering * Julia 实现 https://github.com/JuliaNeighbors/HNSW.jl * Java 实现:https://github.com/jelmerk/hnswlib * 使用 Java Native Access 的 Java 绑定:https://github.com/stepstone-tech/hnswlib-jna * .Net 实现:https://github.com/curiosity-ai/hnsw-sharp * CUDA 实现:https://github.com/js1010/cuhnsw * Rust 实现 https://github.com/rust-cv/hnsw * 出于内存和线程安全目的编写的 Rust 实现,包含一个 Trait 以允许用户实现自定义距离。它接受满足 T:Serialize+Clone+Send+Sync 的类型 T 的切片作为数据:https://github.com/jean-pierreBoth/hnswlib-rs ### 2 亿 (200M) SIFT 测试重现 下载并解压 bigann 数据集(从根目录): ``` python tests/cpp/download_bigann.py ``` 编译: ``` mkdir build cd build cmake .. make all ``` 在 2 亿 (200M) SIFT 子集上运行测试: ``` ./main ``` BigANN 子集的大小(以百万为单位)由硬编码在 **sift_1b.cpp** 中的变量 **subset_size_millions** 控制。 ### 更新测试 生成测试数据(从根目录开始): ``` cd tests/cpp python update_gen_data.py ``` 编译(从根目录开始): ``` mkdir build cd build cmake .. make ``` 运行**不带**更新的测试(从 `build` 目录开始) ``` ./test_updates ``` 运行**带**更新的测试(从 `build` 目录开始) ``` ./test_updates update ``` ### HNSW 示例演示 - 100 万件亚马逊商品的可视化搜索引擎 (MXNet + HNSW):[网站](https://thomasdelteil.github.io/VisualSearch_MXNet/),[代码](https://github.com/ThomasDelteil/VisualSearch_MXNet),由 [@ThomasDelteil](https://github.com/ThomasDelteil) 提供演示 ### 参考文献 HNSW 论文: ``` @article{malkov2018efficient, title={Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs}, author={Malkov, Yu A and Yashunin, Dmitry A}, journal={IEEE transactions on pattern analysis and machine intelligence}, volume={42}, number={4}, pages={824--836}, year={2018}, publisher={IEEE} } ``` 本仓库支持的更新算法将发表在由 Apoorv Sharma、Abhishek Tayal 和 Yury Malkov 撰写的美国专利 "Dynamic Updates For HNSW, Hierarchical Navigable Small World Graphs"(专利号 15/929,802)中。
标签:C++, HNSW算法, Python, 向量检索, 安全意识培训, 数据擦除, 无后门, 机器学习基础库, 近似最近邻, 逆向工具