外部产品的矢量化减少 - NumPy

Question

我对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循环？

提前致谢!

Answer 1

方法#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

第三种方法只是吹响了其他一切!

为什么Approach #3比10x-130x更快Approach #1？

np.einsum在C.实现在第一种方法,与三串有i,j,k在其字符串的符号,我们将有三个嵌套循环(当然,在C).那里有很多内存开销.

随着第三个方法,我们只进入两个字符串i,j,因此两个嵌套循环(在C再次),并借力基于BLAS matrix-multiplication与np.dot.这两个因素是这个因素的惊人加速.