稀疏矩阵上的逐行外积

Question

给定两个稀疏 scipy 矩阵 A, B 我想计算行外积。

我可以通过多种方式使用 numpy 来做到这一点。最简单的可能是

np.einsum('ij,ik->ijk', A, B).reshape(n, -1)

或

(A[:, :, np.newaxis] * B[:, np.newaxis, :]).reshape(n, -1)

其中 n 是 A 和 B 中的行数。

然而，就我而言，遍历密集矩阵会占用太多 RAM。因此，我发现的唯一选择是使用 python 循环：

sp.sparse.vstack((ra.T@rb).reshape(1,-1) for ra, rb in zip(A,B)).tocsr()

虽然使用较少的 RAM，但速度非常慢。

因此，我的问题是，是否有一种稀疏（RAM 高效）方法来获取两个矩阵的行外积，从而保持向量化？

（一个类似的问题是但那里的所有答案都通过密集矩阵。）

Answer 1

我们可以直接计算结果的csr表示。它不是超快（在 100,000x768 上约 3 秒）但可能没问题，具体取决于您的用例：

import numpy as np
import itertools
from scipy import sparse

def spouter(A,B):
    N,L = A.shape
    N,K = B.shape
    drows = zip(*(np.split(x.data,x.indptr[1:-1]) for x in (A,B)))
    data = [np.outer(a,b).ravel() for a,b in drows]
    irows = zip(*(np.split(x.indices,x.indptr[1:-1]) for x in (A,B)))
    indices = [np.ravel_multi_index(np.ix_(a,b),(L,K)).ravel() for a,b in irows]
    indptr = np.fromiter(itertools.chain((0,),map(len,indices)),int).cumsum()
    return sparse.csr_matrix((np.concatenate(data),np.concatenate(indices),indptr),(N,L*K))

A = sparse.random(100,768,0.03).tocsr()
B = sparse.random(100,768,0.03).tocsr()

print(np.all(np.einsum('ij,ik->ijk',A.A,B.A).reshape(100,-1) == spouter(A,B).A))

A = sparse.random(100000,768,0.03).tocsr()
B = sparse.random(100000,768,0.03).tocsr()

from time import time
T = time()
C = spouter(A,B)
print(time()-T)

样本运行：

True
3.1073222160339355

稀疏矩阵上的逐行外积

Row-wise outer product on sparse matrices

python

numpy

matrix

scipy

sparse-matrix