问题描述
我有一个形状=(N,N)的矩阵A和一个形状相同的矩阵B=(N,N)。 我正在使用以下 einsum(使用 opt_einsum 库)构建矩阵 M:
M = oe.contract('nm,in,jm,pn,qm->ijpq',A,B,B)
这是计算以下总和:
这产生了形状为 (N,N,N) 的矩阵 M。然后我将其重塑为形状为 (N**2,N**2) 的二维数组
M = M.reshape((N**2,N**2))
这必须是二维的,因为它被视为线性运算符。
我想使用 sparse 库,因为 M 是稀疏的,并且变得太大而无法存储大 N。我可以使 A 和 B 变得稀疏,然后将它们插入到 oe.contract
中。
问题是,sparse 仅支持 2D 数组,因此无法产生形状 (N,N. N) 的 4D 输出。有没有办法结合 einsum 和 reshape 步骤来允许以这种方式使用稀疏,因为 M 的最终形状是 2D?
解决方法
这可能对您使用 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,N*N)
c3 = np.einsum('n,on,pn->op',A.reshape(N*N),c1)
,
我没有正确解释 opt_einsum
给出的错误。
问题是不是稀疏不支持ND稀疏数组(它支持!),而是我没有使用真正的einsum,因为总和的索引出现了两次以上({{ 1}} 和 n
)。如 m
文档中所述,这将导致使用 sparse.einsum 函数,但该函数不存在。仅使用每个索引中的 1 或 2 个即可。使用不同的路径,例如 hpaulj 建议的路径可以用于解决问题。