使用 HDF5 进行 knn 搜索

Question

我正在尝试对内存有限的大数据进行 knn 搜索。

我正在使用 HDF5 和 python。

我尝试了暴力线性搜索（使用 pytables）和 kd-tree 搜索（使用 sklearn）

令人惊讶，但是 kd-tree 方法需要更多时间（如果我们增加批量大小，kd-tree 可能会更好地工作？但我不知道最佳大小也受内存限制）

现在我正在寻找如何加快计算速度，我认为 HDF5 文件可以针对个人 PC 进行调整，也可以使用 nymexpr 或一些 python 技巧来加快规范计算。

import numpy as np
import time
import tables
import cProfile

from sklearn.neighbors import NearestNeighbors

rows = 10000
cols = 1000
batches = 100
k= 10

#USING HDF5
vec= np.random.rand(1,cols)
data = np.random.rand(rows,cols)
fileName = 'C:\carray1.h5'
shape = (rows*batches, cols)  # predefined size
atom = tables.UInt8Atom()  #?
filters = tables.Filters(complevel=5, complib='zlib') #?

#create
# h5f = tables.open_file(fileName, 'w')
# ca = h5f.create_carray(h5f.root, 'carray', atom, shape, filters=filters)

# for i in range(batches):
    # ca[i*rows:(i+1)*rows]= data[:]+i  # +i to modify data

# h5f.close()

#can be parallel?
def test_bruteforce_knn():
    h5f = tables.open_file(fileName)

    t0= time.time()
    d = np.empty((rows*batches,))
    for i in range(batches):
        d[i*rows:(i+1)*rows] = ((h5f.root.carray[i*rows:(i+1)*rows]-vec)**2).sum(axis=1)
    print (time.time()-t0)
    ndx = d.argsort()
    print ndx[:k]

    h5f.close()

def test_tree_knn():
    h5f = tables.open_file(fileName)

        # it will not work
    # t0= time.time()
    # nbrs = NearestNeighbors(n_neighbors=k, algorithm='ball_tree').fit(h5f.root.carray)
    # distances, indices = nbrs.kneighbors(vec)
    # print (time.time()-t0)

        #need to concatenate distances, indices somehow 
    t0= time.time()
    d = np.empty((rows*batches,))
    for i in range(batches):
        nbrs = NearestNeighbors(n_neighbors=k, algorithm='ball_tree').fit(h5f.root.carray[i*rows:(i+1)*rows])
        distances, indices = nbrs.kneighbors(vec)  # put in dict? 
        #d[i*rows:(i+1)*rows] = 
    print (time.time()-t0)
    #ndx = d.argsort()
    #print ndx[:k]

    h5f.close()

cProfile.run('test_bruteforce_knn()')
cProfile.run('test_tree_knn()')

Answer 1

如果我理解正确，您的数据有 1000 个维度？如果是这种情况，那么预计 kd-tree 不会很好，因为它会遭受维度灾难。

您可能想看看 Approximate Nearest Neighbors 搜索方法。例如看看flann。