如何使用python加速最近邻搜索?

时间:2016-09-23 12:38:07

标签: python performance parallel-processing nearest-neighbor kdtree

我有一个代码,它计算最近的体素(未分配)到体素(已分配)。那就是我有一系列体素,很少有体素已经分配了标量(1,2,3,4 ......等),而且很少有体素是空的(比如值'0')。下面的代码找到最近的指定体素到未指定的体素,并为该体素指定相同的标量。因此,标量为“0”的体素将根据最近的体素分配一个值(1或2或3,...)。以下代码有效,但需要花费太多时间。 有替代方案吗?或者如果您对如何进一步改进它有任何反馈意见?

“”“#self.voxels是一个3D numpy数组”“”

def fill_empty_voxel1(self,argx, argy, argz):
""" where # argx, argy, argz are the voxel location where the voxel is zero"""
    argx1, argy1, argz1 = np.where(self.voxels!=0)   # find the non zero voxels
    a = np.column_stack((argx1, argy1, argz1)) 
    b = np.column_stack((argx, argy, argz))
    tree = cKDTree(a, leafsize=a.shape[0]+1)
    distances, ndx = tree.query(b, k=1, distance_upper_bound= self.mean) # self.mean is a mean radius search value
    argx2, argy2, argz2 = a[ndx][:][:,0],a[ndx][:][:,1],a[ndx][:][:,2]
    self.voxels[argx,argy,argz] = self.voxels[argx2,argy2,argz2] # update the voxel array

实施例

“”“这是一个小数据集的小例子:”“”

import numpy as np
from scipy.spatial import cKDTree
import timeit

voxels = np.zeros((10,10,5), dtype=np.uint8)
voxels[1:2,:,:] = 5.
voxels[5:6,:,:] = 2.
voxels[:,3:4,:] = 1.
voxels[:,8:9,:] = 4.
argx, argy, argz = np.where(voxels==0)

tic=timeit.default_timer()
argx1, argy1, argz1 = np.where(voxels!=0)   # non zero voxels
a = np.column_stack((argx1, argy1, argz1)) 
b = np.column_stack((argx, argy, argz))
tree = cKDTree(a, leafsize=a.shape[0]+1)
distances, ndx = tree.query(b, k=1, distance_upper_bound= 5.)
argx2, argy2, argz2 = a[ndx][:][:,0],a[ndx][:][:,1],a[ndx][:][:,2]
voxels[argx,argy,argz] = voxels[argx2,argy2,argz2]
toc=timeit.default_timer()
timetaken = toc - tic #elapsed time in seconds
print '\nTime to fill empty voxels', timetaken

用于可视化:

from mayavi import mlab
data = voxels.astype('float')
scalar_field = mlab.pipeline.scalar_field(data)
iso_surf = mlab.pipeline.iso_surface(scalar_field)
surf = mlab.pipeline.surface(scalar_field)  
vol = mlab.pipeline.volume(scalar_field,vmin=0,vmax=data.max())  
mlab.outline()    
mlab.show()    

现在,如果我将体素数组的维度设置为(500,500,500),则计算最近搜索所需的时间不再有效。我怎么能克服这个?并行计算可以减少时间(我不知道我是否可以并行化代码,如果你这样做,请告诉我)?

潜在的修复:

通过在cKDTree查询中添加n_jobs = -1参数,我可以大大缩短计算时间。

distances, ndx = tree.query(b, k=1, distance_upper_bound= 5., n_jobs=-1)

对于13核CPU上的(400,100,100)数组,我能够在不到一个小时的时间内计算距离。我尝试使用1个处理器,完成相同的阵列大约需要18个小时。 感谢@gsamaras的答案!

2 个答案:

答案 0 :(得分:1)

尝试sklearn.neighbors.NearestNeighbors会很有趣,它提供n_jobs参数:

  

为邻居搜索运行的并行作业的数量。

这个软件包还提供了Ball Tree算法,你可以测试它与kd-tree的算法,但是我的预感是kd-tree会更好(但这又取决于你的数据,所以研究一下!)

您可能还想使用降维,这很容易。我们的想法是减少您的维度,因此您的数据包含的信息较少,因此可以更快地解决最近邻问题。当然,这里有一个权衡,准确性!

您可能/将会降低维度的准确性,但值得尝试。但是,这通常适用于高维空间,您只是在3D 。所以我不知道是否符合你的特定情况使用sklearn.decomposition.PCA是有意义的。

一句话:

如果你真的想要高性能,你不会得到,你可以切换到,然后使用CGAL

答案 1 :(得分:1)

您可以切换到近似最近邻(ANN)算法,这些算法通常利用复杂的散列或邻近图技术快速索引数据并执行更快的查询。一个例子是Spotify的Annoy。 Annoy的自述文件包括this plot,其中显示了近年来发布的各种ANN算法的精确性能权衡比较。性能最佳的算法hnswNon-Metric Space Library (NMSLIB)下有一个Python实现。