从我的一个earlier posts开始,我注意到操作np.add.at(A, indices, B)
比A[indices] += B
慢得多。
def fast(A):
n = A.shape[0]
retval = np.zeros(2*n-1)
for i in range(n):
retval[slice(i,i+n)] += A[i, :]
return retval
def slow(A):
n = A.shape[0]
retval = np.zeros(2*n-1)
for i in range(n):
np.add.at(retval, slice(i,i+n), A[i, :])
return retval
def slower(A):
n = A.shape[0]
retval = np.zeros(2*n-1)
indices = np.arange(n)
indices = indices[:,None] + indices
np.add.at(retval, indices, A) # bottleneck here
return retval
我的计时:
A = np.random.randn(10000, 10000)
timeit(lambda: fast(A), number=10) # 0.8852798199995959
timeit(lambda: slow(A), number=10) # 56.633683917999406
timeit(lambda: slower(A), number=10) # 57.763389584000834
显然,使用__iadd__
要快得多。但是,np.add.at
的文档指出:
对由“indices”指定的元素的操作数“a”执行无缓冲到位操作。对于加法ufunc,此方法等效于a[indices] += B,除了为索引超过一次的元素累积结果。
为什么np.add.at
这么慢?
此功能的用例是什么?
2条答案
按热度按时间daolsyd01#
add.at
用于索引包含重复项且+=
未生成所需结果的情况add.at
:与显式迭代的结果相同:
一些时间:
add.at
比+=
慢,但比python迭代好。我们可以测试
A[:4]+1
、A[:4]+=1
等变体。无论如何,如果你不需要,不要使用
add.at
变体。编辑
你的例子,简化了一点:
因此,您正在向重叠切片添加值。我们可以用数组替换切片:
不需要添加您的'慢'情况,
add.at
与切片,因为索引没有重复。创建所有索引。
+=
由于缓冲而无法工作add.at
解决了:完整的Python迭代:
现在一些时间。
基线:
高级索引会减慢一些:
如果它工作,一个高级索引+=将是最快的:
完整的python迭代与循环数组的情况大致相同:
add.at
比平坦的+=慢,但仍然比基线好:在我的小型测试中,
slower
的性能最好。但它可能无法像base/fast那样扩展。你的indices
要大得多。通常对于非常大的数组,由于减少了内存管理负载,一维上的迭代速度更快。aij0ehis2#
我也有
np.add.at
速度慢的问题。最后,我基于排序编写了自己的版本:对于小型数组,这比
np.add.at
慢,但对于大型数组,它快20倍或更多。小基准:
大型基准:
还有一种常见的模式,它比
np.add.at
更简单更快,尽管它仍然比排序方法慢: