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

我发现了一种更快的方法，只在组合上循环：

^{pr2}$

有没有更有效的方法？在

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

Answer 1

方法1

我们可以直接使用迭代器作为^{} string notation和NumPy的内置^{}。因此，解决方案是使用一个einsum调用-

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

方法2

我们可以使用broadcasted elementwise-multiplication和np.tensordot或{}来表示所有的{}。在

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

^{pr2}$

然后，tensordot或{}-

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的方法快。在