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)

行是索引数组，其值在范围（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]

可能涉及替代实施切片[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秒的时间内完成。

Answer 2

一种可能性就是简单地对einsum方法进行分块。将rows和cols切成20个大小的小块可以解决笔记本电脑上约2分钟内的大（10 ^ 7）问题。人们可以通过调整块大小来改善这一点。

但我们可以做得更好：我们可以按行或列分组（我选择了cols），然后将各个cols与所有成对的行相乘。我们可以使用稀疏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

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

2 个答案: