组合 Einsum 表达式

Question

我想评价一下

E = np.einsum('ij,jk,kl->ijkl',A,A,A)
F = np.einsum('ijki->ijk',E)

其中 A 是一个矩阵（大小不超过 1000 x 1000）。计算 E 很慢。我想通过只计算我存储在 F 中的“对角线”元素来加快这个速度。是否可以将这两个表达式结合起来？/有没有更好的方法来加快这个计算？

Answer 1

我不确定是否有自动的方法，但您始终可以自己计算并给出最终表达式 einsum：

F = np.einsum('ij,jk,ki->ijk', A, A, A)

Answer 2

In [86]: A=np.random.randint(0,100,(100,100))

In [88]: E1=np.einsum('ijki->ijk',np.einsum('ij,jk,kl->ijkl',A,A,A))

In [89]: E2=np.einsum('ij,jk,ki->ijk',A,A,A)

In [90]: np.allclose(E1,E2)
Out[90]: True

良好的时间改进 - 100 倍，对应于保存的维度 (l)

In [91]: timeit np.einsum('ijki->ijk',np.einsum('ij,jk,kl->ijkl',A,A,A))
1 loops, best of 3: 1.1 s per loop

In [92]: timeit np.einsum('ij,jk,ki->ijk',A,A,A)
100 loops, best of 3: 10.9 ms per loop

einsum 对所有索引执行组合迭代，尽管是在 Cython 代码中。因此，减少索引的数量可以节省大量时间。看起来i...i 组合在初始计算中有效。

只有 2g 的内存，(1000,1000) 太大了 E1 情况下的“迭代器太大”，E2 情况下的“内存错误”。