不使用循环进行多个numpy点积

Question

是否有可能在没有循环的情况下计算多个点积？假设您有以下内容：

a = randn(100, 3, 3)
b = randn(100, 3, 3)

我希望获得一个形状为 (100, 3, 3) 的数组 z，以便对于所有的 i

z[i, ...] == dot(a[i, ...], b[i, ...])

换句话说，这意味着要进行验证：

for va, vb, vz in izip(a, b, z):
    assert (vq == dot(va, vb)).all()

直接的解决方案是：

z = array([dot(va, vb) for va, vb in zip(a, b)])

这里使用了隐式循环（列表推导式+数组）来计算。

是否有更高效的方法来计算z？

- user3716510

---

在Python 3.5中，它将是a @ b（除非自PEP 465以来有所改变）。不幸的是，dot不起作用。我不知道为什么。 - user2357112

Answer 1

np.einsum可以在这里发挥作用。尝试运行此可复制粘贴代码：

import numpy as np

a = np.random.randn(100, 3, 3)
b = np.random.randn(100, 3, 3)

z = np.einsum("ijk, ikl -> ijl", a, b)

z2 = np.array([ai.dot(bi) for ai, bi in zip(a, b)])

assert (z == z2).all()

einsum是编译代码，运行速度非常快，甚至比np.tensordot更快（虽然在这里并不完全适用，但通常适用）。以下是一些统计数据：

In [8]: %timeit z = np.einsum("ijk, ikl -> ijl", a, b)
10000 loops, best of 3: 105 us per loop


In [9]: %timeit z2 = np.array([ai.dot(bi) for ai, bi in zip(a, b)])
1000 loops, best of 3: 1.06 ms per loop

- eickenberg

---

tensordot 在这里非常适用；只需在 axes=[[2],[1]] 上进行收缩。虽然我怀疑 einsum 会更快，考虑到执行的大量小收缩，但我认为 tensordot 更适合在大维度上进行收缩优化。 - Eelco Hoogendoorn

Eelco，使用这些轴进行求和并不能得到与einsum相同的结果。 - Oliver W.

np.tensordot会执行一种类似外积的操作，将a中的每个矩阵与b中的每个矩阵相乘，从而得到2个大小为100的轴。期望的结果是对角线，因此np.tensordot计算出了太多的系数。 - eickenberg

Answer 2

尝试在numpy中使用Einstein求和：

z = np.einsum('...ij,...jk->...ik', a, b)

这个方法优雅且不需要你编写循环，正如你所要求的那样。 在我的系统上，它使速度提高了4.8倍：

%timeit z = array([dot(va, vb) for va, vb in zip(a, b)])
1000 loops, best of 3: 454 µs per loop

%timeit z = np.einsum('...ij,...jk->...ik', a, b)
10000 loops, best of 3: 94.6 µs per loop

- Oliver W.

Answer 3

除了其他答案，我想要补充一点：

np.einsum("ijk, ijk -> ij", a, b)

这适用于我遇到的一个相关案例，其中您有两个由匹配的2D向量（点或方向）的2D字段组成的3D数组。这为这些2D向量之间提供了一种“逐元素”点积。

例如：

np.einsum("ijk, ijk -> ij", [[[1,2],[3,4]]], [[[5,6],[7,8]]])
# => array([[17, 53]])

在哪里：

np.dot([1,2],[5,6])
# => 17
np.dot([3,4],[7,8])
# => 53

- Wacov

Answer 4

这个解决方案仍然使用循环，但是更快，因为它通过使用dot的out参数避免了不必要的临时数组创建：

def dotloop(a,b):
    res = empty(a.shape)
    for ai,bi,resi in zip(a,b,res):
        np.dot(ai, bi, out = resi)
    return res

%timeit dotloop(a,b)
1000 loops, best of 3: 453 us per loop
%timeit array([dot(va, vb) for va, vb in zip(a, b)])
1000 loops, best of 3: 843 us per loop

- shx2