首页 文章

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

提问于
浏览
3

我想以一种在内存使用和计算时间方面有效的方式计算矩阵乘积的某些特定值 . 问题是中间矩阵具有两个非常大的尺寸并且可能不可存储 .

带示例值的尺寸:

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]

另一种实现可能涉及切片[rows]和b [:,cols],创建具有形状(L,K)和(K,L)的矩阵,但这些矩阵也太大了 . Numpy在进行花式切片时复制值

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

提前致谢

2 回答

  • 2

    一种可能性是直接计算结果 . 也许还有一些其他技巧可以使用BLAS例程而没有庞大的临时数组,但这也有效 .

    Example

    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秒的时间内运行 .

  • 1

    一种可能性就是简单地对你的方法进行分块 . 将 rowscols 分成大小为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
    

相关问题