python - 没有循环的多个 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？

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