没有循环的多个numpy dot产品

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)])

它使用隐式循环(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

此解决方案仍然使用循环，但速度更快，因为它通过使用outarg避免了不必要的临时数组创建dot：

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