两个巨大的稠密矩阵的乘法Hadamard--乘以一个稀疏矩阵

Question

我有两个密集矩阵A和B，每个矩阵的大小都是3e5x100。另一个大小为3e5x3e5的稀疏二进制矩阵C。我想要找到以下数量：C ∘ (AB')，其中∘是Hadamard乘积(即元素智慧)，B'是B的转置。显式计算AB'会占用大量内存(~500 of )。由于最终结果不需要整个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语言编写代码，所以我认为对于LAPACK/BLAS来说，numpy可以成为更人性化的包装器。

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