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, 向量检索, 安全意识培训, 数据擦除, 无后门, 机器学习基础库, 近似最近邻, 逆向工具