优化Python：大数组，内存问题

Question

我在运行一个Python/NumPy代码时遇到了速度问题。我不知道如何使它更快，也许有其他人可以帮忙？

假设有一个表面，有两个三角剖分，一个细的（..._fine）有M个点，一个粗的有N个点。此外，每个点上都有粗网格上的数据（N个浮点数）。我正在尝试做以下事情：

对于细网格上的每个点，找到粗网格上最接近的k个点并获取平均值。简单来说，就是从粗到细进行数据插值。

我的代码现在是这样的。对于大数据（在我的情况下M=2e6，N=1e4），代码运行大约需要25分钟，猜测是由于显式的for循环没有进入NumPy。有什么想法用智能索引解决这个问题吗？MxN数组会导致内存溢出..

import numpy as np

p_fine.shape => m x 3
p.shape => n x 3

data_fine = np.empty((m,))
for i, ps in enumerate(p_fine):
    data_fine[i] = np.mean(data_coarse[np.argsort(np.linalg.norm(ps-p,axis=1))[:k]])

干杯！

- Max

---

你不能使用sklearn中的最近邻回归吗？这样可能比手动实现更有效率。 - benten

我认为Numpy不是做这种事情的好模块，因为针对细网格点的循环无法向量化。如果您需要手动编码，我建议使用Cython并使用显式的for循环进行操作。 - Bertrand Gazanion

如果我理解正确的话，p 和 p_fine 是网格。由于网格通常是有结构的，如果您切换到不同的数据结构（例如 kD 树），在其中搜索空间数据会更快。 - Hannes Ovrén

Answer 1

首先感谢详细的帮助。

Divakar，你的解决方案大大加快了速度。对于我的数据，根据块大小的差异，代码运行时间在2分钟左右。

我也尝试过使用sklearn，最终得出了以下结果：

def sklearnSearch_v3(p, p_fine, k):
    neigh = NearestNeighbors(k)
    neigh.fit(p)
    return data_coarse[neigh.kneighbors(p_fine)[1]].mean(axis=1)

对于我的数据大小，最终运行速度相当快，我得到了以下结果：

import numpy as np
from sklearn.neighbors import NearestNeighbors

m,n = 2000000,20000
p_fine = np.random.rand(m,3)
p = np.random.rand(n,3)
data_coarse = np.random.rand(n)
k = 3

产量

%timeit sklearv3(p, p_fine, k)
1 loop, best of 3: 7.46 s per loop

- Max

---

这似乎是更好的选择！你在研究中做得很好。 - Divakar

Answer 2

方法一

我们正在处理大型数据集，并且内存是一个问题，因此我将尝试优化循环中的计算。现在，我们可以使用np.einsum来替换np.linalg.norm部分，并使用np.argpartition替代实际的np.argsort排序，代码如下 -

out = np.empty((m,))
for i, ps in enumerate(p_fine):
    subs = ps-p
    sq_dists = np.einsum('ij,ij->i',subs,subs)
    out[i] = data_coarse[np.argpartition(sq_dists,k)[:k]].sum()
out = out/k

方法 #2

现在，我们还可以采用另一种方法，使用Scipy的cdist函数来实现完全向量化的解决方案，代码如下 -

from scipy.spatial.distance import cdist
out = data_coarse[np.argpartition(cdist(p_fine,p),k,axis=1)[:,:k]].mean(1)

但是，由于我们在这里受到内存限制，因此我们可以分块执行这些操作。基本上，我们将从具有数百万行的高数组 p_fine 中获取行块，并使用 cdist，因此在每次迭代中获取输出元素的块而不仅仅是一个标量。通过这种方式，我们将通过该块的长度减少循环计数。

因此，最终我们将拥有以下实现 -

out = np.empty((m,))
L = 10 # Length of chunk (to be used as a param)
num_iter = m//L
for j in range(num_iter):
    p_fine_slice = p_fine[L*j:L*j+L]
    out[L*j:L*j+L] = data_coarse[np.argpartition(cdist\
                           (p_fine_slice,p),k,axis=1)[:,:k]].mean(1)

运行时测试

安装 -

# Setup inputs
m,n = 20000,100
p_fine = np.random.rand(m,3)
p = np.random.rand(n,3)
data_coarse = np.random.rand(n)
k = 5

def original_approach(p,p_fine,m,n,k):
    data_fine = np.empty((m,))
    for i, ps in enumerate(p_fine):
        data_fine[i] = np.mean(data_coarse[np.argsort(np.linalg.norm\
                                                 (ps-p,axis=1))[:k]])
    return data_fine

def proposed_approach(p,p_fine,m,n,k):    
    out = np.empty((m,))
    for i, ps in enumerate(p_fine):
        subs = ps-p
        sq_dists = np.einsum('ij,ij->i',subs,subs)
        out[i] = data_coarse[np.argpartition(sq_dists,k)[:k]].sum()
    return out/k

def proposed_approach_v2(p,p_fine,m,n,k,len_per_iter):
    L = len_per_iter
    out = np.empty((m,))    
    num_iter = m//L
    for j in range(num_iter):
        p_fine_slice = p_fine[L*j:L*j+L]
        out[L*j:L*j+L] = data_coarse[np.argpartition(cdist\
                               (p_fine_slice,p),k,axis=1)[:,:k]].sum(1)
    return out/k

时间 -

In [134]: %timeit original_approach(p,p_fine,m,n,k)
1 loops, best of 3: 1.1 s per loop

In [135]: %timeit proposed_approach(p,p_fine,m,n,k)
1 loops, best of 3: 539 ms per loop

In [136]: %timeit proposed_approach_v2(p,p_fine,m,n,k,len_per_iter=100)
10 loops, best of 3: 63.2 ms per loop

In [137]: %timeit proposed_approach_v2(p,p_fine,m,n,k,len_per_iter=1000)
10 loops, best of 3: 53.1 ms per loop

In [138]: %timeit proposed_approach_v2(p,p_fine,m,n,k,len_per_iter=2000)
10 loops, best of 3: 63.8 ms per loop

因此，第一种提出的方法有约2倍的改进，第二种方法在len_per_iter参数设置为1000的最佳位置上比原始方法快20倍。希望这将把您的25分钟运行时间缩短到一分钟左右。我想这还不错！

- Divakar