带外积的矢量化减和 - 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 比 Approach #1 快 10-130 倍？

np.einsum 是在 C 中实现的。在第一种方法中，这三个字符串在其字符串表示法中有 i,j,k，我们将有三个嵌套循环（当然是在 C 中） .那里有很多内存开销。

使用第三种方法，我们只进入两个字符串i、j，因此有两个嵌套循环（再次在 C 中），并且还利用基于 BLAS 的 matrix-multiplication 和 np.dot。这两个因素促成了这一惊人的加速。