大规模向量相似度计算(一)——hnswlib的基本使用示例

随着推荐系统基于向量相似的召回策略越发普及和成熟。

推荐系统基于向量相似的召回更加广泛的应用到推荐系统当中。

今天我们来聊一聊如何基于hnswlib来进行大规模的物品和物品之间的相似度计算。
这一篇内容主要讲一下列一下hnswlib使用一个基本demo，具体的参数解析参考：

《大规模向量相似度计算(一)——hnswlib的基本使用示例》
《大规模向量相似度计算(二)——hnswlib的参数含义》
《HNSW的基本原理及使用》

1. 安装 hnswlib

pip install hnswlib

2. 构建索引

import hnswlib
index = hnswlib.Index(space='l2', dim=dim)
index.init_index(max_elements=num_elements, ef_construction=200, M=16)

3. 添加向量

index.add_items(vectors, labels)

4 近邻检索

index.set_ef(int(k * 1.2))
labels, distances = index.knn_query(data, k=k)

5 完整代码

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import hnswlib
import argparse
import requests
import logging
import rediscluster as rc
from contextlib import closing
from retry import retry
import numpy as np

logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s %(levelname)s %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)


def get_image_vector(path, batch_size):
    index = 0
    image_labels = []
    image_vectors = []
    for line in open(path, "r"):
        line = line.strip().split("\t")
        label, vector = line[0], line[1]
        image_labels.append(label)
        image_vectors.append([float(dim) for dim in vector.split(",")])
        index += 1
        if index >= batch_size:
            yield np.array(image_labels), np.array(image_vectors)
            image_labels.clear()
            image_vectors.clear()
            index = 0
    yield np.array(image_labels), np.array(image_vectors)


def build_hnsw_index(path, dim, num_elements, batch_size):
    count = 0
    index = hnswlib.Index(space='l2', dim=dim)
    index.init_index(max_elements=num_elements, ef_construction=200, M=16)
    for labels, vectors in get_image_vector(path, batch_size):
        index.add_items(vectors, labels)
        count += 1
        logging.info("add items index:{}".format(count * batch_size))
    return index


def top_k(index, data, k):
    index.set_ef(int(k * 1.2))
    labels, distances = index.knn_query(data, k=k)
    return labels, distances


def main():
    path = args["input"]
    dim = args["dim"]
    num_elements = args["size"]
    batch_size = args["batch_size"]
    k = args["k"]
    output_path = args["output"]
    output = open(output_path, "w")
    index = build_hnsw_index(path, dim, num_elements, batch_size)
    count = 0
    for labels, vectors in get_image_vector(path, batch_size):
        targets, distances = top_k(index, vectors, k)
        for label, label_targets, label_distances in zip(labels, targets, distances):
            all_targets = []
            for target, distance in zip(label_targets, label_distances):
                line = "{}:{}".format(target, distance)
                all_targets.append(line)

            line = "{}\t{}\n".format(label, ",".join(all_targets))
            output.write(line)
        count += 1
        logging.info("build top k index:{}".format(count * batch_size))


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", default="", type=str, help="path of input items vector")
    ap.add_argument("--dim", default=2048, type=int, help="dim of items vector")
    ap.add_argument("--size", default=6000000, type=int, help="elements number of the items")
    ap.add_argument("--batch_size", default=100000, type=int, help="batch size to process")
    ap.add_argument("--k", default=100, type=int, help="top k result")
    ap.add_argument("--output", default="", type=str, help="output path of the result")
    args = vars(ap.parse_args())
    main()

6 测试数据和代码下载

hnswlib 完整代码

搜索并关注微信公众号：【charlie_mouse】
回复：hnswlib 获取测试代码