優化Python：大數組，內存問題

Question

我遇到了運行python/numypy代碼的速度問題。我不知道如何讓它更快，也許是別人？

假設有一個表面有兩個三角形，一個有M點的罰款（..._罰款），一個有N個點的罰款。另外，每個點都有關於粗網格的數據（N個浮點數）。我正在嘗試執行以下操作：

對於細網格上的每個點，找到粗網格上的k個最近點並獲取平均值。短：內插數據從粗到細。

我現在的代碼就是這樣。對於大數據（在我的情況下，M = 2e6，N = 1e4）代碼運行大約25分鐘，猜測由於明確的循環不會進入numpy。任何想法如何用智能索引來解決這個問題？ M×N陣列吹RAM ..

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

乾杯！

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 

Answer 2

方法＃1

我們正在與大型數據集和內存的工作是一個問題，所以我會盡量在循環中優化計算。現在，我們可以使用np.einsum與np.argsort更換到位實際排序的np.linalg.norm一部分np.argpartition，像這樣 -

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's cdist爲一個完全量化的解決方案，像這樣 -

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

但是，由於我們這裏的內存限制，我們可以執行這些化經營ns大塊。基本上，我們將從具有數百萬行的高排列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 

所以，有大約2x改進與第一提出的方法和20x在與甜蜜點與len_per_iter參數組在1000第二個原來的做法。希望這會使您的25分鐘運行時間稍微減少一分鐘。不錯，我猜！