Numpy：更快地计算涉及总和减少的三重嵌套循环

Question

我的目标是有效地计算以下嵌套循环，

Ab = np.random.randn(1000, 100)    
Tb = np.zeros((100, 100, 100))

for i in range(d):
    for j in range(d):
        for k in range(d):
            Tb[i, j, k] = np.sum(Ab[:, i] * Ab[:, j] * Ab[:, k])

我找到了一种更快的方法来执行嵌套循环，即只循环组合：

for i,j,k in itertools.combinations_with_replacement(np.arange(100), 3):
    Abijk = np.sum(Ab[:, i] * Ab[:, j] * Ab[:, k])

    Tb[i, j, k] = Abijk
    Tb[i, k, j] = Abijk

    Tb[j, i, k] = Abijk
    Tb[j, k, i] = Abijk

    Tb[k, j, i] = Abijk
    Tb[k, i, j] = Abijk

有没有更有效的方法？

我希望有一种方法可以利用 Numpy 的 Blas、Numba 的 JIT 或 Pytorch GPU 实现。

Answer 1

方法 #1

我们可以像 einsum string notation with NumPy's built-in np.einsum 那样直接使用迭代器。因此，解决方案是使用单个 einsum 调用 -

Tb = np.einsum('ai,aj,ak->ijk',Ab,Ab,Ab)

方法 #2

我们可以使用 broadcasted elementwise-multiplication 然后 np.tensordot 或 np.matmul 的组合来获得所有 sum-reductions.

因此，再次使用 einsum 或显式维度扩展和 broadcasting -

获得广播的逐元素乘法

parte1 = np.einsum('ai,aj->aij',Ab,Ab)
parte1 = (Ab[:,None,:]*Ab[:,:,None]

然后，tensordot或np.matmul-

Tb = np.tensordot(parte1,Ab,axes=((0),(0)))
Tb = np.matmul(parte1.T, Ab) # Or parte1.T @ Ab on Python 3.x

因此，第二种方法总共有四种可能的变体。

运行时测试

In [140]: d = 100
     ...: m = 1000
     ...: Ab = np.random.randn(m,d)

In [148]: %%timeit  # original faster method
     ...: d = 100
     ...: Tb = np.zeros((d,d,d))
     ...: for i,j,k in itertools.combinations_with_replacement(np.arange(100), 3):
     ...:     Abijk = np.sum(Ab[:, i] * Ab[:, j] * Ab[:, k])
     ...: 
     ...:     Tb[i, j, k] = Abijk
     ...:     Tb[i, k, j] = Abijk
     ...: 
     ...:     Tb[j, i, k] = Abijk
     ...:     Tb[j, k, i] = Abijk
     ...: 
     ...:     Tb[k, j, i] = Abijk
     ...:     Tb[k, i, j] = Abijk
1 loop, best of 3: 2.08 s per loop

In [141]: %timeit np.einsum('ai,aj,ak->ijk',Ab,Ab,Ab)
1 loop, best of 3: 3.08 s per loop

In [142]: %timeit np.tensordot(np.einsum('ai,aj->aij',Ab,Ab),Ab,axes=((0),(0)))
     ...: %timeit np.tensordot(Ab[:,None,:]*Ab[:,:,None],Ab,axes=((0),(0)))
     ...: %timeit np.matmul(np.einsum('ai,aj->ija',Ab,Ab), Ab)
     ...: %timeit np.matmul(Ab.T[None,:,:]*Ab.T[:,None,:], Ab)


10 loops, best of 3:  56.8 ms per loop
10 loops, best of 3:  59.2 ms per loop
 1 loop,  best of 3: 673   ms per loop
 1 loop,  best of 3: 670   ms per loop

最快的似乎是基于 tensordot 的。因此，35x+ 比基于更快的单循环 itertools 的方法提速。

Numpy：更快地计算涉及总和减少的三重嵌套循环

Numpy: Faster computation of triple nested loops involving sum reductions

python

parallel-processing

jit

gpu

numpy

运行时测试