KNN过时了！ANNs比它快了整整380倍-蒲公英云

全文共5889字，预计学习时长15分钟

KNN过时了！ANNs比它快了整整380倍

图源：Google

我们正在经历一场灭绝性的大事件，颇受欢迎的KNN算法正面临淘汰，而几乎每门数据科学课上都会学习这种算法！

KNN背景

寻找与给定项目的K个相似项的做法在机器学习界被广泛称为“相似”搜索或“最近邻”（NN）搜索。最广为人知的最近邻搜索算法便是K最近邻（KNN）算法。

它的用途很广，在已有的物品集如手机电商目录中，运用KNN,便可以从这一整个目录中找到一个少数（Ｋ）最近邻来发起新的搜索请求。

比如说，在下面这个例子中，如果将K设为３，那每一个“iphone”的３最近邻便是其他“iphone”。类似地，对每一个“Samsung”的3最近邻便是所有的Samsungs。

KNN过时了！ANNs比它快了整整380倍

6款手机的目录

KNN的问题

尽管KNN很适合寻找类似项,但是它使用的是穷尽匹配距离计算方法。如果数据中有1000项的话，要找到一个新产品中K=3的最近邻，该算法就会把新产品与数据库中其他产品一起执行1000次距离计算。

这还不算太糟糕。但试想在现实中，顾客对顾客（C2C）的市场数据库里有着上百万的产品，并且每天都可能会上传新的产品。把新产品与所有数百万的产品进行比对的做法确实太浪费时间了，也就是说根本无法拓展。

解决方法

让最近邻在大量的数据中也适用的解决方法就是彻底规避这种暴力的距离计算法，代之以更复杂的一类算法，名为近似最近邻（ANN）。

KNN过时了！ANNs比它快了整整380倍

近似最近邻（ANN）

严格来说，近似最近邻这种算法在最近邻搜索中允许少量的错误。但在现实中的C2C市场里，“真正的”最近邻的数字要比被搜索的“K”最近邻高。相比暴力的KNN，ANN能够在短时间内达到惊人的准确率。下列有几种ANN算法：

· Spotify的【ANNOY】

· Google的【ScaNN】

· Facebook的【Faiss】

· 还有个人最爱：分层可导航小世界图【HNSW】

下面我们把焦点从Python的 sklearn中的KNN算法转向在Python的hnswlib 包中的HNSW图这一出色的ANN算法。接下来将使用大型的【Amazon product dataset】，其中包含‘手机&配件’分类中的527000个产品，以此来证明HNSW的速度非常快（准确说是快380倍），同时还能得到与sklearn的KNN百分之99.3相同的结果。

在HNSW中【paper @ arxiv】，作者使用多层图来描述ANN算法。在插入元素的时候，HNSW图是通过随机选择每个元素的最大层，以指数递减的概率分布逐步建立的。这保证了layer=0时有很多元素来进行精确搜索，而layer=2时有e^-2较少的元素，便于粗略搜索。

最近邻的搜索从最顶层开始粗略搜索，继而向下层递进，直至在最下层使用贪心算法的线路来遍历全图，找到所需数字的近邻。

KNN过时了！ANNs比它快了整整380倍

HNSW图的结构。在最顶层开始最近邻的搜索（粗略搜索），在最底层结束搜索（精细搜索）。

HNSWPython包

整个HNSW算法都是通过C++写成的，并与Python绑定,能用Python包管理工具(pip)，通过打字安装在你的机器里：pip install hnswlib。安装这个包并导入后，创建HNSW图需要几个步骤，将其包装成下列的便捷函数。

import hnswlib
import numpy as npdef fit_hnsw_index(features, ef=100, M=16,save_index_file=False):
    # Convenience function to create HNSWgraph
    # features : list of lists containingthe embeddings
    # ef, M: parameters to tune the HNSWalgorithm
    num_elements = len(features)
    labels_index =np.arange(num_elements)    EMBEDDING_SIZE= len(features[0])    # Declaring index
    # possible space options are l2,cosine or ip
    p = hnswlib.Index(space='l2',dim=EMBEDDING_SIZE)    # Initing index -the maximum number of elements should be known
   p.init_index(max_elements=num_elements, ef_construction=ef, M=M)    # Element insertion
    int_labels = p.add_items(features,labels_index)    # Controlling the recallby setting ef
    # ef should always be > k
    p.set_ef(ef) 
    # If you want to save the graph to afile
    if save_index_file:
         p.save_index(save_index_file)
    return p

创建HNSW索引后，查询“K”最近邻就像调用下列一行代码一样简单。

ann_neighbor_indices, ann_distances = p.knn_query(features, k)

KNNvs. ANN 基准实验

首先下载一个有50万行以上的大型数据集。接着用预先训练好的 fasttext句子向量，将文本列转换为 300d的嵌入向量。

然后训练KNN和HNSW ANN模型，输入长短不一的数据[1000, 10000,100000, len(data)]，以此来测量数据大小对速度的影响。

最后，从两个模型中寻求K=10和100的最近邻来测量K对速度的影响。首先引入必要的包和模型。这会花一点时间，fasttext模型需要从网上下载。

# Imports
# For input data pre-processing
import json
import gzip
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import fasttext.util
fasttext.util.download_model('en', if_exists='ignore') # English pre-trainedmodel
ft = fasttext.load_model('cc.en.300.bin')# For KNN vs ANN benchmarking
from datetime import datetime
from tqdm import tqdm
from sklearn.neighbors import NearestNeighbors
import hnswlib

数据

使用【Amazon product dataset】，它包含了‘手机&配件’分类中527000个产品。从这个链接下载数据集，并运行以下代码来将其转换成一个数据框架。只需用到产品的title列，因为要用它来寻找相似产品。

# Data: http://deepyeti.ucsd.edu/jianmo/amazon/
data = []
with gzip.open('meta_Cell_Phones_and_Accessories.json.gz') as f:
    for l in f:
       data.append(json.loads(l.strip()))# Pre-Processing: https://colab.research.google.com/drive/1Zv6MARGQcrBbLHyjPVVMZVnRWsRnVMpV#scrollTo=LgWrDtZ94w89
# Convert list into pandas dataframe
df = pd.DataFrame.from_dict(data)
df.fillna('', inplace=True)# Filter unformatted rows
df = df[~df.title.str.contains('getTime')]# Restrict to just 'Cell Phones andAccessories'
df = df[df['main_cat']=='Cell Phones & Accessories']# Reset index
df.reset_index(inplace=True, drop=True)# Only keep the title columns
df = df[['title']]# Check the df
print(df.shape)
df.head()

如果一切运行顺利，就能得到如下输出结果。

KNN过时了！ANNs比它快了整整380倍

亚马逊产品数据集

嵌入

要在文本数据上运行相似性搜索，就必须要首先将其转变为数字向量。一个快捷的方法就是使用预先训练好的网络嵌入层，比如Facebook提供的【FastText】。因为所有行有要有相同长度的向量，因此不用考虑title里的字数，应该在df中将 get_sentence_vector方法运用到title列上。

在嵌入完成后，将emb列作为一个列表提取出来，输入到NN算法中。当然在这一步之前需要先进行文本清理预处理。另外，利用微调后的嵌入模型也是个不错的选择。

# Title Embedding using FastText Sentence Embedding
df['emb'] = df['title'].apply(ft.get_sentence_vector)# Extract out theembeddings column as a list of lists for input to our NN algos
X = [item.tolist() for item in df['emb'].values]

基准

有了算法输入后，就可以开始基准测试了。以搜索空间中的产品数量和被搜索的K最近邻为循环，在循环中运行测试。

在每次迭代时，除了对每个算法所花费的时间进行计时外，还要检查pct_overlap ，作为KNN最近邻和同时被ANN抓取的最近邻数量之比。

注意！整个测试要在8核，30GB RAM的机器上全天候运行六天左右，所以会花一些时间。当然也可以通过多重处理来加速，因为每一次运行实际上是互相独立的。

# Number of products for benchmark loop
n_products = [1000, 10000, 100000, len(X)]# Number of neighbors for benchmarkloop
n_neighbors = [10, 100]# Dictionary to save metric results for each iteration
metrics = {'products':[], 'k':[], 'knn_time':[], 'ann_time':[],'pct_overlap':[]}for products in tqdm(n_products):
    # "products" number ofproducts included in the search space
    features = X[:products]
    for k in tqdm(n_neighbors):   
        # "K" Nearest Neighborsearch
        # KNN 
        knn_start = datetime.now()
        nbrs = NearestNeighbors(n_neighbors=k,metric='euclidean').fit(features)
        knn_distances,knn_neighbor_indices = nbrs.kneighbors(X)
        knn_end = datetime.now()
       metrics['knn_time'].append((knn_end - knn_start).total_seconds())
        # HNSW ANN
        ann_start = datetime.now()
        p = fit_hnsw_index(features,ef=k*10)
        ann_neighbor_indices,ann_distances = p.knn_query(features, k)
        ann_end = datetime.now()
       metrics['ann_time'].append((ann_end - ann_start).total_seconds())
        # Average Percent Overlap inNearest Neighbors across all "products"
       metrics['pct_overlap'].append(np.mean([len(np.intersect1d(knn_neighbor_indices[i],ann_neighbor_indices[i]))/k for i in range(len(features))]))
        metrics['products'].append(products)
        metrics['k'].append(k)
metrics_df = pd.DataFrame(metrics)
metrics_df.to_csv('metrics_df.csv', index=False)
metrics_df