Skip to content

SJTU-DMTai/HEXA

Repository files navigation

HEXA

Compile

Step:

  1. 修改CmakeLists.txt, 把nvcc的路径正确写入
set(CMAKE_CUDA_COMPILER /path/to/nvcc)
  1. 填写下列参数中的路径/path/to/xxx,然后运行
mkdir cmake-build-release
cmake -DCMAKE_BUILD_TYPE=Release -DCMAKE_MAKE_PROGRAM=/path/to/ninja -G Ninja -DPython_EXECUTABLE=/path/to/python -DPython3_EXECUTABLE=/path/to/python -S . -B ./cmake-build-release
cmake --build cmake-build-release --target hello_world

data preparation

base vector 与query vector数据格式,二进制文件用于存储 n 个 dim 维度向量,整体结构:

  1. 前 8 个字节为两个小端存储的 int32 类型数据,依次是向量总数n、单个向量维度dim;
  2. 后续紧跟$n\times dim\times sizeof(Type)$个字节数据,采用小端存储、行主序排列,完整存储 n 个 dim 维向量;
  3. 向量数据类型Type支持三种:float32(4 字节 / 元素)、uint8(1 字节 / 元素)、int8(1 字节 / 元素),其中行主序表示先完整存储第 1 个向量的所有维度,再依次存储后续向量。

ground truth数据格式,二进制文件存储id与距离

  1. 前 8 个字节为两个小端存储的 int32 类型数据,依次是查询数目nq、每个查询的groundtruth数目K
  2. 后续$4\times nq\times K$ 采用小端存储、行主序排列,uint32格式,完整存储 nq 个查询的近邻id
  3. 后续$sizeof(DistType) \times nq\times K$ 采用小端存储、行主序排列,完整存储 nq 个查询的近邻距离query的距离

数据源:

Dataset Base Query Groundtruth Vector Type Distance Type
BigANN1B BigANN1B Base BigANN1B Query 10K BigANN1B GT uint8 int32
DEEP1B DEEP1B Base DEEPP1B Query 10K DEEP1B GT float float
SPACEV1B SPACEV1B Base SPACEV1B Query SPACEV1B GT int8 int
SSNPP1B SSNPP1B base SSNPP1B Query 100K - uint8 int32

注意:

  1. ground truth文件中需要提供距离,主要是考虑到一些数据集上会出现距离相同排名并列的情况,需要基于距离给出更加精确的recall。如果使用上述数据源,BigANN、DEEP、SPACEV均已经给出了距离。如果使用其他不包含距离的数据源,我们给出了python代码gt_to_gt_with_dist.py来从仅包含id的ground truth文件生成我们所需的包含距离的格式。
  2. 我们强烈建议数据集规模至少为100-million scale
  3. SSNPP数据集的groundtruth需要自己手动计算,为此我们提供了多GPU加速的计算groundtruth的代码,如下
mkdir cmake-build-release
cmake -DCMAKE_BUILD_TYPE=Release -DCMAKE_MAKE_PROGRAM=/path/to/ninja -G Ninja -DPython_EXECUTABLE=/path/to/python -DPython3_EXECUTABLE=/path/to/python -S . -B ./cmake-build-release
cmake --build cmake-build-release --target make_gt
./cmake-build-release/src/make_gt --base /path/to/ssnpp/base --query /path/to/ssnpp/query --output /tmp_file.bin --dtype uint8 --dist l2 --gpus <gpu list, e.g. 0,1,2,3> --batch_size 1000000 
# modify path in gt_to_gt_with_dist.py 
python gt_to_gt_with_dist.py 

build index

final_build build \
  --data <data.bin> \
  --index <index.bin> \
  --vec-type <float|uint8|int8> \
  --dist-type <l2|cosine> \
  --n-graph <int> \
  --ele-per-postlist <int> \
  --out-degree <int> \
  [--dev-list 0,1,2,3] \
  [--limit-n 0]

参数说明:

  • data:基向量文件
  • index:输出索引文件
  • vec-type:向量类型 (float,uint8,int8)
  • dist-type:距离类型
  • n-graph:subgraph数量(预期,实际可能会小于该值)
  • ele-per-postlist:每个second-level cluster中元素数目
  • out-degree:图出度
  • dev-list:使用的 GPU 设备列表
  • limit-n:限制参与构建的向量数量,0 表示全部,大于0表示只是用前若干个向量

示例

/home/xuyifei/projects/bi_index/cmake-build-release/src/final_build build \
    --data /data1/xuyifei/ann_data/sift100m/base.fbin \
    --index /data1/xuyifei/ann_data/bi_index_data/final_sift100m/index1.bin \
    --vec-type float \
    --dist-type l2 \
    --n-graph 80 \
    --ele-per-postlist 50000 \
    --out-degree 64 \
    --dev-list 0,1,2,3

Compression(Optional)

final_build compress \
  --index <index.bin> \
  --out <index_cmp.bin> \
  [--dev-list 0,1,2,3]

示例:

cmake-build-release/src/final_build compress \
  --index /data1/xuyifei/ann_data/bi_index_data/final_sift100m/index1.bin \
  --out /data1/xuyifei/ann_data/bi_index_data/final_sift100m/index1_cmp.bin \
  --dev-list 0,1,2,3

Search and evaluation

final_build search \
  --data <data.bin> \
  --index <index.bin> \
  --query <query.bin> \
  --gt <groundtruth.bin> \
  --vec-type <float|uint8|int8> \
  --dist-type <l2|cosine> \
  --triple <K,efSearch,glook[,reach_limit]> \
  --triple <K,efSearch,glook[,reach_limit]> ...

说明:

  • triple 支持多次传入
  • reach_limit 可省略,省略时为 0
  • 输出包含 QPS 与召回率评估
  • 无论index是build后直接产生的index,还是compression产生的index,都是用相同的命令进行search
  • efSearch 对应论文中的B,lookGraph对应论文中$K^g$,reachLimit对应论文中R
  • efSearch不得小于K

示例:

ERR_FILE="/home/xuyifei/projects/bi_index/log/final/sift100m_search_${TIMESTAMP}.err.log"

echo -e "searching using original index: \n" | tee -a "$LOG_FILE"

/home/xuyifei/projects/bi_index/cmake-build-release/src/final_build search \
    --data /data1/xuyifei/ann_data/sift100m/base.fbin \
    --index /data1/xuyifei/ann_data/bi_index_data/final_sift100m/index1.bin \
    --query /data1/xuyifei/ann_data/sift100m/query.fbin \
    --gt /data1/xuyifei/ann_data/sift100m/groundtruth.ibin \
    --vec-type float \
    --dist-type l2 \
    --triple 10,200,16,4 \
    --triple 10,20,2,3\
    --triple 10,20,3,3\
    --triple 10,30,4,4\
    --triple 10,40,5,4\
    --triple 10,70,8,5\
    --triple 10,100,10,6\
    --triple 10,200,16,7\
    --triple 10,400,16,8\
    2> >(tee -a "$ERR_FILE" >&2) | tee -a "$LOG_FILE"

Configuration recommendation (optional)

final_build optimize \
  --data <data.bin> \
  --index <index.bin> \
  --query <query.bin> \
  --gt <groundtruth.bin> \
  --vec-type <float|uint8|int8> \
  --dist-type <l2|cosine> \
  --param --K <K> --efSearch <list> --lookGraph <list> --reachLimit <list> 

说明:

  • --param 用于分隔不同参数组, efSearch 对应论文中的B,lookGraph对应论文中$K^g$,reachLimit对应论文中R
  • 列表参数表示搜索空间,使用逗号分隔,例如 --efSearch 10,20,30
  • 无论index是否被compression都可以用相同的命令调用
  • 这里的参数推荐基于你提供的合成的query,与真实query之间存在差异;建议把所有推荐的参数配置增加一档后使用

示例

/home/xuyifei/projects/bi_index/cmake-build-release/src/final_build optimize \
    --data /data1/xuyifei/ann_data/sift100m/base.fbin \
    --index /data1/xuyifei/ann_data/bi_index_data/final_sift100m/index1.bin \
    --query /data1/xuyifei/ann_data/sift100m/sampled_query.10k.fbin \
    --gt /data1/xuyifei/ann_data/sift100m/sampled_groundtruth.10k.bin \
    --vec-type float \
    --dist-type l2 \
    --param --K 10 --efSearch 1,3,5,10,20,30,40,50,60,80,100,150,200,300,400,500,600,700,800,900,1000 --lookGraph 1,2,3,4,5,6,7,8,10,12,14,16,18,20,23,26,30,35 --reachLimit 1,2,3,4,5,6,7,8,9,10,12,14,16

About

Implementation for vldb paper, HEXA: A Disjoint-Subgraph-Based Indexing Framework for Approximate Nearest Neighbor Search at Billion Scale

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages