python - 结合sparse和einsum进行大稀疏求和

Question

我有一个形状=（N，N）的矩阵A和一个形状相同的矩阵B=（N，N）。我正在使用以下 einsum（使用opt_einsum库）构建矩阵 M：

M = oe.contract('nm,in,jm,pn,qm->ijpq', A, B, B, B, B)

这是评估以下总和：

这会产生形状为 (N, N, N, N) 的矩阵 M。然后我将其重塑为形状的二维数组 (N**2, N**2)

M = M.reshape((N**2, N**2))

这必须是二维的，因为它被视为线性算子。

我想使用稀疏库，因为 M 是稀疏的，并且变得太大而无法存储大 N。我可以使 A 和 B 稀疏，并将它们插入到oe.contract.

问题是，sparse 仅支持 2D 数组，因此无法生成形状 (N, N, N. N) 的 4D 输出。有没有办法结合 einsum 和 reshape 步骤以允许以这种方式使用 sparse ，因为 M 的最终形状是 2D ？

标签： pythonarrayssparse-matrixnumpy-einsum

Answer 1

这可能对您的使用没有帮助opt_einsum，但通过一些重组我可以加快np.einsum相当多的速度，至少对于小型数组。

做两个的部分乘积B：

c1 = np.einsum('in,jm->ijnm',B,B).reshape(N*N,N,N)

这pq对是相同的，所以我们不需要重新计算它：

c2 = np.einsum('nm,onm,pnm->op',A,c1,c1)

我验证这适用于两个 (3,3) 数组，并且速度提高了大约 10 倍。

我们甚至可以将其重塑nm为 1d，尽管这并不能提高速度：

c1 = np.einsum('in,jm->ijnm',B,B).reshape(N*N,N*N)
c3 = np.einsum('n,on,pn->op',A.reshape(N*N),c1,c1)