优化Python：大型数组、内存问题

我在运行python/numypy代码时遇到速度问题。我不知道如何让它更快，也许是其他人

假设有一个曲面有两个三角剖分，一个精细（…_精细）有M个点，一个粗糙有N个点。此外，每个点（N个浮点数）都有粗网格上的数据。我正在努力做到以下几点：

对于细网格上的每个点，在粗网格上找到k个最近点并获得平均值。短：从粗到细插值数据

我现在的代码是这样的。对于大数据（在我的例子中，M=2e6，N=1e4），由于显式for循环没有进入numpy，所以代码运行大约25分钟。有没有办法用智能索引解决这个问题？M x N阵列吹扫闸板

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]])

干杯

方法#1

我们正在处理大型数据集，内存是一个问题，因此我将尝试优化循环中的计算。现在，我们可以使用替换

np.linalg.norm

部分，并用

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

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

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)

方法#2

现在，作为另一种方法，我们也可以用于完全矢量化的解决方案，如下所示-

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

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

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)

但是，由于这里内存有限，我们可以分块执行这些操作。基本上，我们将从拥有数百万行的高数组

p_fine

中获取行块，并使用

cdist

，从而在每次迭代中获得输出元素块，而不仅仅是一个标量。有了这个，我们就可以将循环计数减少该块的长度

最后，我们将有一个这样的实现-

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

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

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

---

因此，第一种方法比原始方法提高了大约

2x

，第二种方法比原始方法提高了

20x

，len_per_iter参数设置为

1000

。希望这将使您的25分钟运行时间缩短到一分钟多一点。我想还不错

首先感谢您的详细帮助

首先，Divakar，您的解决方案大大加快了速度。根据我的数据，代码运行了不到2分钟，这取决于块大小

我也试着用我的方法来学习，结果是

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

---

有什么原因不能在sklearn中使用？我认为numpy不是做这类事情的好模块，因为精细网格点上的循环无法矢量化。如果您需要手工编写代码，我建议使用Cython和explicit for循环。如果我理解正确，

p

和

p\u fine

是网格。由于网格通常是结构化的，因此如果切换到搜索空间数据速度较快的其他数据结构（例如kD树），则搜索速度会快得多。这似乎是更好的方法！在研究这些方面做得很好。