Question

是否可以在没有循环的情况下计算多个点积？说你有以下内容：

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

我希望得到一个形状为z的数组{100}，以便所有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)])

使用隐式循环（list comprehension + array）。

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

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

Answer 2

尝试在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 3

此解决方案仍然使用循环，但速度更快，因为它使用out的{{1}} arg避免了不必要的临时数组创建：

dot

Answer 4

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

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

适用于我遇到的相关案例，其中有两个3D阵列，包括2D矢量（点或方向）的匹配2D场。这给了一种＆＃34;元素明智的＆＃34;这些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

没有循环的多个numpy dot产品

4 个答案: