NumPy 中节省时间和空间的数组乘法

Question

给定 NumPy 数组 R 和 S 的形状分别为 (m, d) 和 (m, n, d)，我想计算一个形状为 (m, n) 的数组 P，其 (i, j)-th 条目是 np.dot(R[i, :] , S[i, j, :]) .

执行双 for 循环不需要任何额外空间（P 的 m * n 空间除外），但不会节省时间。

使用广播，我可以使用P = np.sum(R[:, np.newaxis, :] * S, axis=2)，但这需要额外的m * n * d 空间。

什么是最节省时间和空间的方法？

Answer 1

einsum 是另一个常见的嫌疑人

m, n, d = 100, 100, 100
>>> R = np.random.random((m, d))
>>> S = np.random.random((m, n, d))
>>> np.einsum('md,mnd->mn', R, S)

>>> np.allclose(np.einsum('md,mnd->mn', R, S), (R[:,None,:]*S).sum(axis=-1))
True
>>> from timeit import repeat
>>> repeat('np.einsum("md,mnd->mn", R, S)', globals=globals(), number=1000)
[0.7004671019967645, 0.6925274690147489, 0.6952172230230644]
>>> repeat('(R[:,None,:]*S).sum(axis=-1)', globals=globals(), number=1000)
[3.0512512560235336, 3.0466731210472062, 3.044075728044845]

一些间接证据表明einsum 不会太浪费内存：

>>> m, n, d = 1000, 1001, 1002
>>> # Too much for broadcasting:
>>> np.zeros((m, n, d))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
MemoryError
>>> R = np.random.random((m, d))
>>> S = np.random.random((n, d))
>>> np.einsum('md,nd->mn', R, S).shape
(1000, 1001)

Answer 2

在这些情况下，最好考虑numba，它可以提供两全其美的效果：

import numpy as np
from numba import jit

def vanilla_mult(R, S):
    m, n = R.shape[0], S.shape[1]
    result = np.empty((m, n), dtype=R.dtype)
    for i in range(m):
        for j in range(n):
            result[i, j] = np.dot(R[i, :], S[i, j,:])
    return result

def broadcast_mult(R, S):
    return np.sum(R[:, np.newaxis, :] * S, axis=2)

@jit(nopython=True)
def jit_mult(R, S):
    m, n = R.shape[0], S.shape[1]
    result = np.empty((m, n), dtype=R.dtype)
    for i in range(m):
        for j in range(n):
            result[i, j] = np.dot(R[i, :], S[i, j,:])
    return result

注意，vanilla_mult 和 jit_mult 具有完全相同的实现，但是后者是即时编译的。让我们测试一下：

In [1]: import test # the above is in test.py

In [2]: import numpy as np

In [3]: m, n, d = 100, 100, 100

In [4]: R = np.random.rand(m, d)

In [5]: S = np.random.rand(m, n, d)

好的...

In [6]: %timeit test.broadcast_mult(R, S)
100 loops, best of 3: 1.95 ms per loop

In [7]: %timeit test.vanilla_mult(R, S)
100 loops, best of 3: 11.7 ms per loop

哎呀，与广播相比，计算时间增加了近 5 倍。不过……

In [8]: %timeit test.jit_mult(R, S)
The slowest run took 760.57 times longer than the fastest. This could mean that an intermediate result is being cached.
1 loop, best of 3: 870 µs per loop

不错！我们可以通过简单的 JITing 将运行时间减半！这个规模如何？

In [12]: m, n, d = 1000, 1000, 100

In [13]: R = np.random.rand(m, d)

In [14]: S = np.random.rand(m, n, d)

In [15]: %timeit test.vanilla_mult(R, S)
1 loop, best of 3: 1.22 s per loop

In [16]: %timeit test.broadcast_mult(R, S)
1 loop, best of 3: 666 ms per loop

In [17]: %timeit test.jit_mult(R, S)
The slowest run took 7.59 times longer than the fastest. This could mean that an intermediate result is being cached.
1 loop, best of 3: 83.6 ms per loop

扩展性非常好，因为广播开始因必须创建大型中间数组而受到阻碍，与普通方法相比，它的时间只有一半，但几乎是 7 倍与 JIT 方法一样多！

编辑添加

最后，我们比较np.einsum 方法：

In [19]: %timeit np.einsum('md,mnd->mn', R, S)
10 loops, best of 3: 59.5 ms per loop

而且它显然是速度的赢家。不过，我对它还不够熟悉，无法评论空间要求。