如何在Python中高效地计算矩阵乘积的稀疏值？

Question

我希望能够高效地计算矩阵乘积的一些特定值，既要考虑内存使用效率，也要考虑计算时间。但问题在于，中间矩阵具有两个非常大的维度，可能无法存储。

以下是示例值的维度：

N = 7  # very large
K = 3
M = 10 # very large
L = 8  # very very large

'a'是一个形状为(N,K)的矩阵
'b'是一个形状为(K,N)的矩阵

a = np.arange(N*K).reshape(N,K)
b = np.arange(K*M).reshape(K,M)

rows是一个包含在范围(N)内且长度为L的索引数组
cols是一个包含在范围(M)内且长度为L的索引数组

rows = [0,0,1,2,3,3,4,6]
cols = [0,9,5,8,2,8,3,6]

我需要以下内容，但由于其大小无法计算中间结果为形状为(MxN)的矩阵(a @ b):

values = (a @ b)[rows, cols]

另一种实现方法可能涉及对a[rows]和b[:,cols]进行切片，创建形状为(L,K)和(K,L)的矩阵，但这些也太大了。 在进行高级切片时，Numpy会复制这些值。

values = np.einsum("ij,ji->i", a[rows], b[:,cols])

提前感谢您的帮助。

- checkThisOut

---

请问您能否提供K、N、M、L的实际数字（大致范围）？ - Paul Panzer

[行数，列数]相对于整个输出大概有多大？只有少数值，5％，50％吗？ - max9111

目前我有 K=1e2, N=1e4, M=1e4, L=1e6，但我希望我的算法能够处理十倍的因子 K=1e3, N=1e5, M=1e5, L=1e7。 - checkThisOut

无法计算 (a @ b)，但是可以计算 np.dot(a[np.unique(rows),:], b[:,np.unique(cols)]) 吗？当然，这在很大程度上取决于您的 rows 和 cols 向量的性质。 - BenBoulderite

我已经删除了所有没有值的行和列，因此unique在这种情况下没有帮助。 - checkThisOut

Answer 1

一种可能性是直接计算结果。也许有其他技巧可以使用BLAS例程而不需要一个巨大的临时数组，但这也可以起作用。

示例

import numpy as np
import numba as nb
import time


@nb.njit(fastmath=True,parallel=True)
def sparse_mult(a,b_Trans,inds):
  res=np.empty(inds.shape[0],dtype=a.dtype)

  for x in nb.prange(inds.shape[0]):
    i=inds[x,0]
    j=inds[x,1]
    sum=0.
    for k in range(a.shape[1]):
      sum+=a[i,k]*b_Trans[j,k]
    res[x]=sum
  return res


#-------------------------------------------------
K=int(1e3)
N=int(1e5)
M=int(1e5)
L=int(1e7)

a = np.arange(N*K).reshape(N,K).astype(np.float64)
b = np.arange(K*M).reshape(K,M).astype(np.float64)

inds=np.empty((L,2),dtype=np.uint64)
inds[:,0] = np.random.randint(low=0,high=N,size=L) #rows
inds[:,1] = np.random.randint(low=0,high=M,size=L) #cols

#prepare
#-----------------------------------------------
#sort inds for better cache usage
inds=inds[np.argsort(inds[:,1]),:]

# transpose b for easy SIMD-usage
# we wan't a real transpose here not a view
b_T=np.copy(np.transpose(b))

#calculate results
values=sparse_mult(a,b_T,inds)

计算步骤，包括准备工作（排序，b矩阵的转置）应该在60秒内完成。

- max9111

---

打算接下来尝试Numba。好主意。 - hilberts_drinking_problem

Answer 2

一种可能的方法是简单地将您的 einsum 方法分块。将 rows 和 cols 切成大小为20的块可在我的笔记本电脑上解决大型（10 ^ 7）问题，大约需要2分钟。通过调整块的大小，可能可以进一步改善。

但我们可以做得更好：我们可以按行或列进行分组（我选择了列），然后将单个列与所有成对行相乘。我们可以使用稀疏的csc / csr矩阵来为我们完成所有的排序/洗牌/重新索引。在相同的数据上使用这种方法只需约30秒即可完成。

import numpy as np
from scipy import sparse

def f_sparse_helper(a, b, rows, cols):
    h = sparse.csr_matrix((np.empty(L), cols, np.arange(L+1)), (L, M)) \
              .tocsc()
    for i in range(M):
        l, r = h.indptr[i:i+2]
        h.data[l:r] = a[rows[h.indices[l:r]]] @ b[:, i]
    return h.tocsr().data

def f_chunk(a, b, rows, cols, chunk=20):
    out = np.empty(L)
    for j in range(0, rows.size, chunk):
        l = j+chunk
        out[j:l] = np.einsum("ij,ji->i", a[rows[j:l]], b[:,cols[j:l]])
    return out

def prep_data(K, M, N, L):
    a = np.random.uniform(0, 10, (N, K))
    b = np.random.uniform(0, 10, (K, M))
    rows = np.random.randint(0, N, (L,))
    cols = np.random.randint(0, M, (L,))
    return a, b, rows, cols

# use small exmpl to check correct
K, M, N, L = 10, 100, 100, 1000
a, b, rows, cols = prep_data(K, M, N, L)
res = f_sparse_helper(a, b, rows, cols)
assert np.allclose(res, np.einsum("ij,ji->i", a[rows], b[:,cols]))
assert np.allclose(res, f_chunk(a, b, rows, cols))

# timeit on big one
from time import perf_counter as pc
K, M, N, L = 1_000, 10_000, 10_000, 10_000_000
a, b, rows, cols = prep_data(K, M, N, L)
t = pc()
res_ch = f_chunk(a, b, rows, cols)
s = pc()
print('chunked      ', s-t, 'seconds')
t = pc()
res_sh = f_sparse_helper(a, b, rows, cols)
s = pc()
print('sparse helper', s-t, 'seconds')
assert np.allclose(res_sh, res_ch)

样例运行:

chunked       121.16188396583311 seconds
sparse helper 31.172512074932456 seconds

- Paul Panzer