当结果的稀疏性已知时的稀疏矩阵乘法（在 python|scipy|cython 中）

Question

假设我们想为给定的稀疏矩阵 A,B 计算 C=A*B，但对 C 的一个非常小的条目子集感兴趣，由索引对列表表示：
行=[i1, i2, i3 ... ]
cols=[j1, j2, j3 ... ]
A 和 B 都很大（比如 50Kx50K），但非常稀疏（

我们如何计算这个乘法子集？

这是一个运行速度非常慢的简单实现：

def naive(A, B, rows, cols):
    N = len(rows)
    vals = []
    for n in xrange(N):
        v = A.getrow(rows[n]) * B.getcol(cols[n])
        vals.append(v[0, 0])

    R = sps.coo_matrix((np.array(vals), (np.array(rows), np.array(cols))), shape=(A.shape[0], B.shape[1]), dtype=np.float64)
    return R

即使对于小矩阵，这也很糟糕：

import scipy.sparse as sps
import numpy as np
D = 1000

A = np.random.randn(D, D)
A[np.abs(A) > 0.1] = 0
A = sps.csr_matrix(A)
B = np.random.randn(D, D)
B[np.abs(B) > 0.1] = 0
B = sps.csr_matrix(B)

X = np.random.randn(D, D)
X[np.abs(X) > 0.1] = 0
X[X != 0] = 1
X = sps.csr_matrix(X)
rows, cols = X.nonzero()
naive(A, B, rows, cols)

在我的机器上，naive() 在 1 分钟后完成，大部分精力都花在了构造行/列上（在 getrow() 和 getcol() 中）。
当然，将这个（非常小的）示例转换为密集矩阵，计算大约需要 100 毫秒：

A0 = np.array(A.todense())
B0 = np.array(B.todense())
X0 = np.array(X.todense())
A0.dot(B0) * X0

关于如何有效地计算这种矩阵乘法有什么想法吗？

• 注意：此问题与以下问题几乎相同： Subset of a matrix multiplication, fast, and sparse
  但是，A 和 B 是 full 矩阵，并且其中一个维度非常低（例如，10） 提议的解决方案似乎同时受益于两者。

Answer 1

稀疏矩阵的格式在这里很重要。您总是需要 A 的行和 B 的列。因此，将 A 存储为 csr 并将 B 存储为 csc 以消除 getrow/getcol 开销。不幸的是，这只是故事的一小部分。

最佳解决方案很大程度上取决于稀疏矩阵的结构（大量稀疏列/行等），但您可以根据字典和集合尝试一种。对于矩阵A，每行保留以下内容：

• 该行上所有非零列索引的集合
• 以非零索引为键，以相应的非零值作为值的字典

对于矩阵B，每一列都保留了类似的字典和集合。

为了计算乘法结果中的元素（M，N），将A的M行与B的N列相乘。乘法：

• 找到非零集的交集
• 计算非零元素的乘法总和（即上面的交集）

在大多数情况下，这应该非常快，因为在稀疏矩阵中，集合的交集通常非常小。

一些代码：

class rowarray():
    def __init__(self, arr):
        self.rows = []
        for row in arr:
            nonzeros = np.nonzero(row)[0]
            nzvalues = { i: row[i] for i in nonzeros }
            self.rows.append((set(nonzeros), nzvalues))

    def __getitem__(self, key):
        return self.rows[key]

    def __len__(self):
        return len(self.rows)


class colarray(rowarray):
    def __init__(self, arr):
        rowarray.__init__(self, arr.T)


def maybe_less_naive(A, B, rows, cols):
    N = len(rows)
    vals = []
    for n in xrange(N):
        nz1,v1 = A[rows[n]]
        nz2,v2 = B[cols[n]]
        # list of common non-zeros
        nz = nz1.intersection(nz2)
        # sum of non-zeros
        vals.append(sum([ v1[i]*v2[i] for i in nz]))

    R = sps.coo_matrix((np.array(vals), (np.array(rows), np.array(cols))), shape=(len(A), len(B)), dtype=np.float64)
    return R

D = 1000

Ap = np.random.randn(D, D)
Ap[np.abs(Ap) > 0.1] = 0
A = rowarray(Ap)
Bp = np.random.randn(D, D)
Bp[np.abs(Bp) > 0.1] = 0
B = colarray(Bp)

X = np.random.randn(D, D)
X[np.abs(X) > 0.1] = 0
X[X != 0] = 1
X = sps.csr_matrix(X)
rows, cols = X.nonzero()
maybe_less_naive(A, B, rows, cols)

这样效率更高一点，测试的乘法大约需要 2 秒（80 000 个元素）。结果似乎基本相同。

---

几场比赛的表演。

对每个输出元素执行两个操作：

• 设置交点
• 乘法

集合交集的复杂度应该是 O(min(m,n)) 其中 m 和 n 是每个操作数中非零的数量。这与矩阵的大小无关，只有每行/列的平均非零数很重要。

乘法（和 dict 查找）的数量取决于在上面的交集中找到的非零的数量。

如果两个矩阵都以概率（密度）p 随机分布非零，并且行/列长度为 n，则：

• 设置交点：O(np)
• 字典查找，乘法：O(np^2)

这表明，对于真正稀疏的矩阵，找到交点是关键点。这也可以通过分析来验证；大部分时间都花在计算交叉点上。

当这反映在现实世界中时，我们似乎花费了大约 20 us 来处理 80 个非零的行/列。这并不是快得令人眼花缭乱，而且代码当然可以做得更快。 Cython 可能是一种解决方案，但这可能是 Python 不是最佳解决方案的问题之一。用 C 语言编写的排序整数的简单线性匹配（合并排序类型算法）应该至少快一个数量级。

需要注意的重要一点是，该算法可以同时针对多个元素并行完成。无需满足于单个线程，因为计算是独立的，只要一个线程处理一个输出点。