我是 NumPy 的新手,经常读到您应该避免编写循环。在许多情况下,我知道如何处理,但目前我有以下代码:
p = np.arange(15).reshape(5,3)
w = np.random.rand(5)
A = np.sum(w[i] * np.outer(p[i], p[i]) for i in range(len(p)))
有人知道是否有办法避免内部 for 循环吗?
提前致谢!
最佳答案
方法#1:使用np.einsum
-
np.einsum('ij,ik,i->jk',p,p,w)
方法#2:使用broadcasting
+ np.tensordot
-
np.tensordot(p[...,None]*p[:,None], w, axes=((0),(0)))
方法 #3: 使用 np.einsum
+ np.dot
-
np.einsum('ij,i->ji',p,w).dot(p)
运行时测试
第 1 组:
In [653]: p = np.random.rand(50,30)
In [654]: w = np.random.rand(50)
In [655]: %timeit np.einsum('ij,ik,i->jk',p,p,w)
10000 loops, best of 3: 101 µs per loop
In [656]: %timeit np.tensordot(p[...,None]*p[:,None], w, axes=((0),(0)))
10000 loops, best of 3: 124 µs per loop
In [657]: %timeit np.einsum('ij,i->ji',p,w).dot(p)
100000 loops, best of 3: 9.07 µs per loop
第 2 组:
In [658]: p = np.random.rand(500,300)
In [659]: w = np.random.rand(500)
In [660]: %timeit np.einsum('ij,ik,i->jk',p,p,w)
10 loops, best of 3: 139 ms per loop
In [661]: %timeit np.einsum('ij,i->ji',p,w).dot(p)
1000 loops, best of 3: 1.01 ms per loop
第三种方法就搞砸了!
为什么方法 #3
比方法 #1
快 10 到 130 倍?
np.einsum
是用 C 实现的。在第一种方法中,这三个字符串有 i
,j
,k
在其字符串表示法中,我们将有三个嵌套循环(当然在 C 中)。那里有很多内存开销。
使用第三种方法,我们只进入两个字符串 i
,j
,因此有两个嵌套循环(再次在 C 中)并且还利用基于 BLAS 的 矩阵乘法
与 np.dot
。这两个因素导致了这个惊人的加速。
关于python - 带外积的矢量化和约简 - NumPy,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/47954401/