Question

我正在尝试获得以下循环的快速矢量化版本：

for i in xrange(N1):
   A[y[i]] -= B[i,:]

此处A.shape = (N2,N3)，y.shape = (N1) y取值[0,N2[，B.shape = (N1,N3)。您可以将y条目的条目视为A行。此处N1很大，N2非常小，N3很小。

我只想做

A[y] -= B

会起作用，但问题是y中有重复的条目，这样做不对（即如果y=[1,1]那么A[1]只会添加一次，不是两次）。此外，这似乎没有比未向量化的for循环更快。

有更好的方法吗？

编辑：YXD将this answer链接到评论中，这些评论最初似乎符合该法案。看起来你可以做到我想要的

np.subtract.at(A, y, B)

并且它确实有效，但是当我尝试运行它时，显着慢于 。所以，问题仍然存在：是否有更高效的方法来做到这一点？

EDIT2：一个例子，使事情具体化：

n1,n2,n3 = 10000, 10, 500
A = np.random.rand(n2,n3)
y = np.random.randint(n2, size=n1)
B = np.random.rand(n1,n3)

for循环，当在ipython中使用%timeit运行时，在我的机器上提供：

10 loops, best of 3: 19.4 ms per loop

subtract.at版本最终为A生成相同的值，但速度要慢得多：

  1 loops, best of 3: 444 ms per loop

Answer 1

原始基于for循环的方法的代码看起来像这样 -

def for_loop(A):
    N1 = B.shape[0]
    for i in xrange(N1):
       A[y[i]] -= B[i,:]
    return A

案例＃1

如果n2＆gt;＆gt; n3，我建议使用这种矢量化方法 -

def bincount_vectorized(A):

    n3 = A.shape[1]
    nrows = y.max()+1
    id = y[:,None] + nrows*np.arange(n3)
    A[:nrows] -= np.bincount(id.ravel(),B.ravel()).reshape(n3,nrows).T
    return A

运行时测试 -

In [203]: n1,n2,n3 = 10000, 500, 10
     ...: A = np.random.rand(n2,n3)
     ...: y = np.random.randint(n2, size=n1)
     ...: B = np.random.rand(n1,n3)
     ...: 
     ...: # Make copies
     ...: Acopy1 = A.copy()
     ...: Acopy2 = A.copy()
     ...: 

In [204]: %timeit for_loop(Acopy1)
10 loops, best of 3: 19 ms per loop

In [205]: %timeit bincount_vectorized(Acopy2)
1000 loops, best of 3: 779 µs per loop

案例＃2

如果n2＆lt;＆lt; n3，可以建议一种具有较小环路复杂度的改进的for循环方法 -

def for_loop_v2(A):
    n2 = A.shape[0]
    for i in range(n2):
        A[i] -= np.einsum('ij->j',B[y==i]) # OR (B[y==i]).sum(0)
    return A

运行时测试 -

In [206]: n1,n2,n3 = 10000, 10, 500
     ...: A = np.random.rand(n2,n3)
     ...: y = np.random.randint(n2, size=n1)
     ...: B = np.random.rand(n1,n3)
     ...: 
     ...: # Make copies
     ...: Acopy1 = A.copy()
     ...: Acopy2 = A.copy()
     ...: 

In [207]: %timeit for_loop(Acopy1)
10 loops, best of 3: 24.2 ms per loop

In [208]: %timeit for_loop_v2(Acopy2)
10 loops, best of 3: 20.3 ms per loop

将添加矢量化为另一个数组索引的数组

1 个答案: