python - 如何在 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])

提前致谢

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 秒内运行。