Question

我需要计算海量数据中2 xyz点之间的距离（100 Gb，大约20个trylion点）。我想加快这个循环。我创建了KDtree，添加了并行计算，将我的数组拆分为更小的部分。所以我想加速的就是这个循环。我的纯python计算时间大约需要10小时42分钟。增加numpy减少时间为5小时34分钟。添加numba速度可达4小时15分钟。但它仍然不够快。我听说Cython是python计算的最快方式，但我没有任何c经验，我不知道如何将我的函数转换为cython代码。如何使用cython或任何其他方式让这个循环更快地运行？

def controller(point_array, las_point_array):  

    empty = []


    tree = spatial.cKDTree(point_array, leafsize=1000, copy_data = True)   

    empty = __pure_calc(las_point_array, point_array, empty, tree)  

    return ptList   

#############################################################################################

@autojit
def __pure_calc(las_point_array, point_array, empty, tree):

    for i in las_point_array:
            p = tree.query(i)   

            euc_dist = math.sqrt(np.sum((point_array[p[1]]-i)**2))  

            ##add one row at a time to empty list
            empty.append([i[0], i[1], i[2], euc_dist, point_array[p[1]][0], point_array[p[1]][1], point_array[p[1]][2]]) 

    return empty

我附上样本数据进行测试：

Sample

Answer 1

您的函数会构建一个列表（closestPt），最终看起来像这样：

[
    [i0[0], i0[1], i0[2], distM0],
    [i1[0], i1[1], i1[2], distM1],
    ...
]

您应该做的第一件事是将整个结果预先分配为NumPy数组（np.empty()），并一次写入一行。这将避免大量的内存分配。然后您会注意到，您可以将sqrt()推迟到最后，并在完成循环后在distM列上运行它。

如果您发布带有随机/样本输入数据的完整工作测试工具，可能会有更多优化机会。

Answer 2

关键是尽可能多地利用矢量化函数，因为在循环内对纯python函数的任何调用都会或多或少地使autojit变得毫无意义（瓶颈将是纯函数调用）。我注意到查询函数是可矢量化的，欧几里德距离计算也是如此。我不确定你的控制器函数中的ptList变量是什么（示例有点错误），但假设它是你的jit函数的输出，或者关闭它，你应该可以做类似的事情这样：

def controller(point_array, las_point_array):

    tree = spatial.cKDTree(point_array, leafsize=1000, copy_data = True)
    distances, pt_idx = tree.query(las_point_array)
    nearest_pts = point_array[pt_idx]
    euc_distances = np.sqrt((nearest_pts - las_point_array).sum(axis=1) ** 2)
    result = np.vstack((las_point_array.T, euc_distances.T, nearest_pts.T)).T

    return result

加速python 3.5循环以像python一样快速运行它

2 个答案: