python - Numpy 切片速度慢？

Question

嗨，我正在使用 numpy + numba 运行科学计算。我已经意识到 numpy 数组就地添加非常慢......与 matlab 相比

这是matlab代码：

tic;
% A,B are 2-d matrices, ind may not be distinct
for ii=1:N 
    A(ind(ii),:) =  A(ind(ii),:) +  B(ii,:);
end
toc;

这是numpy代码：

s = time.time()
# A,B are numpy.ndarray, ind may not be distinct
for k in xrange(N):
     A[ind[k],:] += B[k,:];
print time.time() - s

结果表明，numpy 代码比 matlab 慢 10 倍……这让我很困惑。

此外，当我将加法从 for 循环中拉出，并将单个矩阵加法与 numpy.add 进行比较时，numpy 和 matlab 的速度似乎相当。

我知道的一个因素是 matlab 使用 JIT for version>=2012a 来加速 for 循环，但我在 python 代码上尝试了 numba，它仍然没有加速。我认为这与 numba 根本没有触及 numpy.add 函数有关，因此性能根本没有改变。

我猜 matlab 为这种情况做了一些病态的缓存，因此它大大击败了 numpy。

关于如何加快 numpy 的任何建议？

score 3 · Accepted Answer

尝试

A[ind] += B[:N]

即没有任何循环。

如果ind可能有重复的元素，您可以使用np.add.at：

np.add.at(A, ind, B[:N])

score 0 · Accepted Answer

这是使用点阵乘法的版本。它构造一个由 1 和 0 组成的矩阵ind。

def bar(A,B,ind):
    K,M =B.shape
    N,M =A.shape
    I = np.zeros((N,K))
    I[ind,np.arange(K)] = 1
    return A+np.dot(I,B)

对于像K,M,N = 30,14,15这样大小的问题，速度大约快 3 倍。但是对于像K,M,N = 300,100,150它这样的大一点来说，它有点慢。

2 回答 2