两个巨大的密集矩阵逐元素相乘，再与一个稀疏矩阵相乘

Question

我有两个密集矩阵A和B，它们的大小都为3e5x100。另外还有一个稀疏二进制矩阵C，大小为3e5x3e5。我想要计算以下值：C ∘ (AB')，其中∘是Hadamard乘积（即逐元素相乘），B'是B的转置。显式计算AB'将需要大量的内存（~500GB）。由于最终结果不需要整个AB'，所以只需计算乘法A_iB_j'，其中C_ij != 0，其中A_i是矩阵A的第i列，C_ij是矩阵C中位置(i,j)的元素。建议的方法如下所示：

result = numpy.initalize_sparse_matrix(shape = C.shape)
while True:
 (i,j) = C_ij.pop_nonzero_index() #prototype function returns the nonzero index and then points to the next nonzero index
 if (i,j) is empty:
   break
 result(i,j) = A_iB_j'

这个算法需要太多时间。有没有办法使用LAPACK/BLAS算法来改进它？我用Python编程，所以我认为numpy可以作为LAPACK/BLAS的更人性化的封装。

- rando

---

你的C语言代码有多稀疏？获取C语言中所有1的坐标需要多长时间？我不记得有任何算法可以减少这个计算。 - Bing Wang

Answer 1

假设C被存储为scipy.sparse矩阵，您可以使用以下方法进行此计算：

C = C.tocoo()
result_data = C.data * (A[C.row] * B[C.col]).sum(1)
result = sparse.coo_matrix((result_data, (row, col)), shape=C.shape)

在这里，我们展示了对于一些较小的输入，结果与朴素算法相匹配：

import numpy as np
from scipy import sparse

N = 300
M = 10

def make_C(N, nnz=1000):
  data = np.random.rand(nnz)
  row = np.random.randint(0, N, nnz)
  col = np.random.randint(0, N, nnz)
  return sparse.coo_matrix((data, (row, col)), shape=(N, N))


A = np.random.rand(N, M)
B = np.random.rand(N, M)
C = make_C(N)

def f_naive(C, A, B):
  return C.multiply(np.dot(A, B.T))

def f_efficient(C, A, B):
  C = C.tocoo()
  result_data = C.data * (A[C.row] * B[C.col]).sum(1)
  return sparse.coo_matrix((result_data, (C.row, C.col)), shape=C.shape)

np.allclose(
    f_naive(C, A, B).toarray(),
    f_efficient(C, A, B).toarray()
)
# True

在这里，我们可以看到它适用于完整的输入大小：

N = 300000
M = 100

A = np.random.rand(N, M)
B = np.random.rand(N, M)
C = make_C(N)

out = f_efficient(C, A, B)

print(out.shape)
# (300000, 300000)

print(out.nnz)
# 1000

- jakevdp