HNSWlib：高效的C++近似最近邻搜索库

icy 昨天 13 抢沙发

默认

摘要： HNSWlib：高效的C++近似最近邻搜索库项目概述 HNSWlib是一个用C++编写的开源库，实现了分层可导航小世界（Hierarchical Navigable Small...

HNSWlib：高效的C++近似最近邻搜索库

项目概述

HNSWlib是一个用C++编写的开源库，实现了分层可导航小世界（Hierarchical Navigable Small World，HNSW）图算法，用于高效的近似最近邻（Approximate Nearest Neighbor，ANN）搜索。该项目由Yury Malkov等人开发，是当前最先进的ANN算法之一，广泛应用于机器学习、信息检索和推荐系统等领域。

核心特性

1. 高性能

极快的搜索速度：相比传统方法，HNSW在保持高召回率的同时，搜索速度提升数个数量级
内存效率：优化的数据结构设计，减少内存占用
多线程支持：支持并行构建索引和搜索

2. 灵活性

多种距离度量：支持欧氏距离（L2）、内积距离、余弦相似度等
动态索引：支持增量添加数据点，无需重建整个索引
可调参数：提供丰富的参数配置，可根据需求平衡精度和速度

3. 易用性

简洁的API：提供直观的C++接口
Python绑定：通过Python接口方便集成到数据科学工作流中
跨平台：支持Linux、macOS和Windows系统

算法原理简介

HNSW算法基于多层图结构，将数据点组织成多个层次的图： 1. 底层图：包含所有数据点，连接密集 2. 上层图：包含部分数据点，连接稀疏，用于快速导航 3. 搜索过程：从上层开始，逐层向下搜索，最终在底层找到最近邻

这种分层结构使得搜索复杂度从O(N)降低到O(log N)，同时保持高召回率。

安装与使用

安装方法

text

# 克隆项目
git clone https://github.com/nmslib/hnswlib.git
cd hnswlib

# 编译
mkdir build
cd build
cmake ..
make -j

# 安装Python绑定
pip install hnswlib

基本使用示例

C++示例

text

#include <hnswlib/hnswlib.h>
#include <vector>
#include <random>
#include <iostream>

int main() {
    // 定义维度
    const int dim = 128;
    const int max_elements = 10000;
    
    // 创建空间（使用L2距离）
    hnswlib::L2Space space(dim);
    
    // 创建HNSW索引
    hnswlib::HierarchicalNSW<float>* alg_hnsw = new hnswlib::HierarchicalNSW<float>(
        &space, max_elements, 16, 200, 100);
    
    // 生成随机数据
    std::mt19937 rng(42);
    std::uniform_real_distribution<> distrib(0.0, 1.0);
    
    // 添加数据点
    for (int i = 0; i < 1000; i++) {
        std::vector<float> point(dim);
        for (int j = 0; j < dim; j++) {
            point[j] = distrib(rng);
        }
        alg_hnsw->addPoint(point.data(), i);
    }
    
    // 搜索最近邻
    std::vector<float> query(dim);
    for (int j = 0; j < dim; j++) {
        query[j] = distrib(rng);
    }
    
    // 执行搜索
    const int k = 5;
    auto result = alg_hnsw->searchKnn(query.data(), k);
    
    // 输出结果
    std::cout << "Top " << k << " neighbors:" << std::endl;
    while (!result.empty()) {
        std::pair<float, hnswlib::labeltype> item = result.top();
        std::cout << "ID: " << item.second << ", Distance: " << item.first << std::endl;
        result.pop();
    }
    
    // 清理资源
    delete alg_hnsw;
    
    return 0;
}

Python示例

text

import hnswlib
import numpy as np

# 创建索引
dim = 128
num_elements = 10000
p = hnswlib.Index(space='l2', dim=dim)

# 初始化索引
p.init_index(max_elements=num_elements, ef_construction=200, M=16)

# 生成随机数据
data = np.float32(np.random.random((num_elements, dim)))

# 添加数据
p.add_items(data)

# 设置查询参数
p.set_ef(50)  # 设置查询时的ef参数

# 生成查询向量
query = np.float32(np.random.random((1, dim)))[0]

# 执行搜索
labels, distances = p.knn_query(query, k=5)

print("Nearest neighbors IDs:", labels)
print("Distances:", distances)

# 保存和加载索引
p.save_index("test_index.bin")
p2 = hnswlib.Index(space='l2', dim=dim)
p2.load_index("test_index.bin")

高级功能

1. 自定义距离度量

text

// 自定义距离函数示例
class CustomDistance : public hnswlib::SpaceInterface<float> {
public:
    size_t data_size_;
    
    CustomDistance(size_t dim) : data_size_(dim * sizeof(float)) {}
    
    size_t get_data_size() override {
        return data_size_;
    }
    
    hnswlib::DISTFUNC<float> get_dist_func() override {
        return custom_dist_func;
    }
    
    void *get_dist_func_param() override {
        return nullptr;
    }
    
    static float custom_dist_func(const void *pVect1, const void *pVect2, const void *qty_ptr) {
        const float *a = (const float *)pVect1;
        const float *b = (const float *)pVect2;
        
        // 实现自定义距离计算
        float result = 0.0f;
        for (size_t i = 0; i < data_size_ / sizeof(float); i++) {
            float diff = a[i] - b[i];
            result += diff * diff;
        }
        return result;
    }
};

2. 批量操作

text

// 批量添加数据点
void batchAddPoints(hnswlib::HierarchicalNSW<float>* index, 
                   const std::vector<std::vector<float>>& data) {
    #pragma omp parallel for
    for (size_t i = 0; i < data.size(); i++) {
        index->addPoint(data[i].data(), i);
    }
}

// 批量搜索
std::vector<std::vector<std::pair<float, hnswlib::labeltype>>> 
batchSearch(hnswlib::HierarchicalNSW<float>* index,
           const std::vector<std::vector<float>>& queries,
           int k) {
    std::vector<std::vector<std::pair<float, hnswlib::labeltype>>> results(queries.size());
    
    #pragma omp parallel for
    for (size_t i = 0; i < queries.size(); i++) {
        auto result = index->searchKnn(queries[i].data(), k);
        std::vector<std::pair<float, hnswlib::labeltype>> neighbors;
        while (!result.empty()) {
            neighbors.push_back(result.top());
            result.pop();
        }
        results[i] = neighbors;
    }
    
    return results;
}

性能优化建议

1. 参数调优

M：每个节点的最大连接数（通常16-64）
efConstruction：构建时的动态候选列表大小（通常100-200）
efSearch：搜索时的动态候选列表大小（通常50-100）

2. 内存优化

text

// 使用内存映射文件处理大数据集
#include <hnswlib/hnswlib.h>

// 创建支持内存映射的索引
hnswlib::HierarchicalNSW<float>* createMMapIndex(
    const std::string& filename,
    hnswlib::SpaceInterface<float>* space,
    size_t max_elements) {
    
    // 使用内存映射优化大索引
    return new hnswlib::HierarchicalNSW<float>(
        space, filename, false, max_elements);
}

实际应用场景

1. 图像检索系统

text

// 简化的图像特征检索示例
class ImageRetrievalSystem {
private:
    hnswlib::HierarchicalNSW<float>* index_;
    hnswlib::L2Space* space_;
    std::unordered_map<int, std::string> id_to_path_;
    
public:
    ImageRetrievalSystem(int feature_dim) {
        space_ = new hnswlib::L2Space(feature_dim);
        index_ = new hnswlib::HierarchicalNSW<float>(
            space_, 1000000, 32, 200, 100);
    }
    
    void addImage(const std::string& path, const std::vector<float>& features) {
        static int next_id = 0;
        index_->addPoint(features.data(), next_id);
        id_to_path_[next_id] = path;
        next_id++;
    }
    
    std::vector<std::string> searchSimilar(const std::vector<float>& query, int k=10) {
        auto result = index_->searchKnn(query.data(), k);
        std::vector<std::string> similar_images;
        
        while (!result.empty()) {
            auto item = result.top();
            similar_images.push_back(id_to_path_[item.second]);
            result.pop();
        }
        
        return similar_images;
    }
};

2. 推荐系统

text

# 基于物品的协同过滤推荐
import hnswlib
import numpy as np
from sklearn.preprocessing import normalize

class ItemBasedCF:
    def __init__(self, num_items, embedding_dim):
        self.index = hnswlib.Index(space='cosine', dim=embedding_dim)
        self.index.init_index(max_elements=num_items, ef_construction=200, M=16)
        self.item_embeddings = {}
        
    def add_item(self, item_id, embedding):
        normalized_embedding = normalize(embedding.reshape(1, -1))[0]
        self.index.add_items(normalized_embedding, item_id)
        self.item_embeddings[item_id] = normalized_embedding
        
    def recommend(self, item_id, top_k=10):
        if item_id not in self.item_embeddings:
            return []
            
        query = self.item_embeddings[item_id]
        labels, distances = self.index.knn_query(query, k=top_k+1)
        
        # 排除自身
        recommendations = []
        for label, distance in zip(labels[0][1:], distances[0][1:]):
            recommendations.append((label, 1 - distance))  # 转换为相似度
            
        return recommendations