NANN (Neural Approximate Nearest Neighbor Search，又名“二向箔”)是一个基于纯TensorFlow的灵活、高性能的大规模检索框架。
NANN是一种任意神经网络相似性度量下的近似近邻检索算法，于2021年在阿里巴巴内被提出并进行了深度优化，后在淘宝展示广告、神马搜索等众多业务上得到了广泛的应用。 NANN对GPU和CPU上的推理性能进行了深度优化，使其能服务于搜索、推荐、广告等对精度和性能有极高要求的场景。 此外，NANN基于原生TensorFlow，易于使用和部署。
简单地说，NANN主要优点可以分为模型训练、性能优化和用户友好性三方面。以下逐一介绍：
- 任意复杂模型：模型训练与索引构建解耦，因此对模型结构几乎没有任何限制，也避免了索引和模型训练绑定带来的高额训练负担。
- 对抗训练：我们采用对抗训练来保证复杂模型下的优越检索性能。
- 高效检索：我们使用 TensorFlow Custom Ops 实现了HNSW 检索过程。就在线检索而言，重写后的 HNSW 检索比 Faiss 原生版本更加高效。
- 运行时优化：我们支持 GPU Multi-Streaming with Multi-Contexts，这极大地增强了并行性。
- 编译优化：我们支持 XLA 并加速了其JIT 过程； 此外，我们还将 XLA 应用到了大规模检索场景，其中batch-size始终是动态的。
- 图级优化：我们针对推荐、搜索和广告领域中常用的一些常见的模型结构，基于 TensorFlow Grappler 提供了一些图级优化。
- 原生TensorFlow：NANN 的后端服务和前端实现完全基于 TensorFlow 生态系统。
- 模型推理和检索解耦：训练-检索解耦使得用户能专注优化深度模型，无需额外考虑检索流程。
- 性能测试：我们提供了一个简单的基准测试工具，可用于分析延迟、吞吐量等推理性能。
git clone [email protected]:alibaba/nann.git docker pull alinann/nann_devel:10.2-cudnn7-devel-ubuntu18.04
| tag | TensorFlow | Python | CUDA | OS | Bazel |
|---|---|---|---|---|---|
| 10.2-cudnn7-devel-ubuntu18.0 | 1.15.5 | 3.7.4 | 10.2 | Ubuntu18.04 | 0.24.1 |
该docker镜像基于 nvidia/cuda:10.2-cudnn7-devel-ubuntu18.04 构建，在/opt/conda路径下安装了TensorFlow runtime和编译的必要依赖，也包含Faiss等相关依赖包。
sudo yum -y install systemd-devel systemd-libs libseccomp device-mapper-libs
sudo mknod /dev/nvidia-modeset c 195 254
DOCKER_PATH=alinann/nann_devel:10.2-cudnn7-devel-ubuntu18.04
## pay attention to the nvidia-driver version, here is 460.73.01
sudo docker run -ti --net=host --volume $HOME:$HOME -w $HOME --volume=nvidia_driver_460.73.01:/usr/local/nvidia:ro --device=/dev/nvidiactl --device=/dev/nvidia-uvm --device=/dev/nvidia-uvm-tools --device=/dev/nvidia0 --device=/dev/nvidia1 --name=tf_whl_open_source_nann $DOCKER_PATH /bin/bashcd tensorflow
./configure # only cuda support is needed
## build whl package for python frontend
bazel build --copt=-mavx2 -c opt --config=cuda --copt -mfpmath=both --copt -mfma --copt -msse4.2 --copt -DGOOGLE_CUDA=1 //tensorflow/tools/pip_package:build_pip_package
./bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg
## compile tf lib for c++ backend and benchmarking
bazel build -c opt --copt -g --strip=never --copt=-mavx --copt=-mavx2 --config=cuda //tensorflow:libtensorflow_framework.so
bazel build -c opt --copt -g --strip=never --copt=-mavx --copt=-mavx2 --config=cuda //tensorflow:libtensorflow_cc.so
## install python frontend
pip install /tmp/tensorflow_pkg/tensorflow-1.15.5-cp37-cp37m-linux_x86_64.whl
# build blaze-benchmark, depend on Tensorflow so files
cd blaze-benchmark
./Build.sh
以下内容对NANN算法的训练、索引结构构建、效果及在线性能评测、部署等流程提供了详细的demo。
本demo基于 UserBehavior 数据集训练和测试，并基于tf.distribute.MirroredStrategy实现了单机多卡并行训练。
- 选项1：下载原始数据集，并使用如下脚本将数据集转为tfrecord格式（大约需要10小时）、训练模型，模型文件将保存在
${output_root}/model目录下：
cd NANN_impls;
# 数据集格式转化
export PYTHONPATH=${PYTHONPATH}:$(pwd);
python nann/data_provider/convert_UB_to_tfrecord.py -i path/to/UserBehavior.csv -o ./data
# 训练模型
eps=3e-5 # adv-eps 控制对抗训练过程中扰动向量的大小，为0时不施加对抗扰动；
batch_size=800
num_neg=200 # 控制采样负样本的数量；
output_root=$(pwd)/output/adv${eps}_bs${batch_size}_neg${num_neg}
python main.py --job-type train --adv-eps ${eps} --num-neg ${num_neg} --batch-size ${batch_size} --output-root ${output_root}- 选项2：为了方便大家走通流程，我们提供了预训练模型和转化后的测试数据集（Google Drive），可以直接使用：
cd NANN_impls;
# 转化后的测试数据集
tar xvf UserBehavior_tfrecords.tar ./
# 预训练模型
output_root=$(pwd)/output/pretrained
mkdir -p ${output_root}
tar xvf nann_ub_ckpt_adv3e-5_bs800_neg200.tar -C ${output_root}模型训练完毕后，我们可以提取Target embedding用于构建HNSW索引和在线部署。 item_ids.npy, item_embs.npy 将保存在 ${output_root}/embeddings 目录下。
python main.py --job-type extract_feature --output-root ${output_root}基于上述target embedding构建HNSW索引，保存在 ${output_root}/index 目录下。
python nann/delivery/build_hnsw_index.py -i ${output_root}/embeddings/item_embs.npy -o ${output_root}/index
测试HNSW检索召回率：
python main.py --job-type test --output-root ${output_root}测试全量打分召回率，用于评估模型打分召回的天花板：
python main.py --job-type test_all --output-root ${output_root}
本阶段生成几个 GraphDef protobufs 用于在线部署：
- exec.pb (二值格式)/ exec.pbtxt (文本格式)：储存模型打分和检索流程。
- frozen_graph.pb：储存模型打分过程，包含模型结构和模型参数。该文件将在
exec.pb中以tf.blaze_xla_op（我们实现的一个custom op）的形式被调用。
cd NANN_impls
export PYTHONPATH=$PYTHONPATH:nann
# export the modified ckpt for online inference,
python main.py --job-type export --output-root ${output_root}
# generate frozen_graph.pb (freeze model and small minor changes).
python nann/delivery/convert_meta.py --model_dir ${output_root}/export --output_dir ${output_root}/export/converted_model
# generate exec.pb. the outer graph (dense part + retrieval) for online inference
# generate exec.meta.pb, only need the signature_def which contains the tensor names of model I/O
python nann/delivery/build_opt_graph.py --model-dir ${output_root}/export --index-dir ${output_root}/index --item-embs-dir ${output_root}/embeddings --output-dir ${output_root}/export可通过以下脚本来验证上述生成文件的正确性：
## exec.meta.pb is only for signature_def, which is used to feed and fetch model I/O
python nann/delivery/NANN_inference_demo.py --graph-file ${output_root}/export/exec.pb --meta-path ${output_root}/export/exec.meta.pb我们提供了一个简单的性能评测工具blaze-benchmark用于分析推理性能：
gen_runmeta.py：用于生成输入模型的mock数据。gen_benchmark_conf.py：用于生成性能评测所需的配置文件。
使用方式如下：
cd NANN_impls
nann_impl_dir=`pwd`
cd nann/benchmark
## gen mock data for benchmarking
python gen_runmeta.py --output-dir ${output_root}/export/
## gen benchmark conf
python gen_benchmark_conf.py --model-dir ${output_root}/export
## open nvidia mps for acceleration, https://docs.nvidia.com/deploy/mps/index.html
/usr/local/nvidia/bin/nvidia-cuda-mps-control -d
## benchmarking
## TF_XLA_CUBIN_CACHE_DIR is the env var for CUBIN cache
## TF_XLA_PTX_CACHE_DIR is the env var for PTX and HLO cache.
## the caches are used for XLA acceleration and reused for the next warmup.
cd ${output_root}/export
TF_CPP_MIN_VLOG_LEVEL=0 TF_XLA_CUBIN_CACHE_DIR=./cubin_cache TF_XLA_PTX_CACHE_DIR=./xla_cache `dirname $nann_impl_dir`/blaze-benchmark/build/benchmark/benchmark benchmark_conf我们基于 TensorFlow Serving 提供了一个简单的demo。
将模型导出为SavedModel格式：
python nann/delivery/pb_to_saved_model.py --model-path ${output_root}/export/exec.pb --meta-path ${output_root}/export/exec.meta.pb --export-dir ${output_root}/export/nann/1基于Docker起模型服务进程：
docker run -p 8501:8501 -p 8500:8500 \
--mount type=bind,source=${output_root}/export/nann,target=/models/nann \
-e MODEL_NAME=nann -t alinann/nann_serving简单的冒烟测试代码如下（tensorflow-serving-api==1.15.0）：
from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2_grpc
import numpy as np
import grpc
import tensorflow as tf
channel = grpc.insecure_channel('localhost:8500')
stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
grpc_request = predict_pb2.PredictRequest()
grpc_request.model_spec.name = "nann"
grpc_request.model_spec.signature_name = 'serving_default'
comm_seq = np.zeros((1, 3200), dtype=np.float16)
level_topn = np.array([100, 200, 200, 200, 200, 200], dtype=np.int32)
grpc_request.inputs['comm_seq'].CopyFrom(tf.make_tensor_proto(comm_seq, shape=comm_seq.shape))
grpc_request.inputs['level_topn'].CopyFrom(tf.make_tensor_proto(level_topn, shape=level_topn.shape))
result = stub.Predict(grpc_request,10)
print(result.outputs['top_k'])article{chen2022approximate,
title={Approximate Nearest Neighbor Search under Neural Similarity Metric for Large-Scale Recommendation},
author={Chen, Rihan and Liu, Bin and Zhu, Han and Wang, Yaoxuan and Li, Qi and Ma, Buting and Hua, Qingbo and Jiang, Jun and Xu, Yunlong and Deng, Hongbo and others},
journal={arXiv preprint arXiv:2202.10226},
year={2022}
}
阿里巴巴已经制定了一套行为准则，我们希望项目参与者能够遵守。 详情请参阅阿里巴巴开源行为准则 (English Version).