无循环进行多个numpy点积

Question

有没有办法在不使用循环的情况下计算多个点积？

假设你有以下内容：

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

我想得到一个形状为 z 的数组，大小是 (100, 3, 3)，这样对于所有的 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 呢？

Answer 1

除了其他回答之外，我想补充一点：

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

这适用于我遇到的一个相关情况，那里有两个3D数组，它们由匹配的2D向量（点或方向）组成。这种情况可以进行一种“逐元素”的点积运算，意思是对这些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

Answer 2

这个解决方案虽然还是用了循环，但速度更快，因为它避免了不必要的临时数组的创建，而是使用了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

Answer 3

试试在numpy中使用爱因斯坦求和法：

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

Answer 4

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