python - Numpy 根据索引求和多维数组中的元素

Question

我正在处理非常大的多维数据，但让我以二维数组为例。给定一个每次迭代都会改变的值数组，

arr = np.array([[ 1, 2, 3, 4, 5], [5, 6, 7, 8, 9]]) # a*b

和一个始终固定的索引数组。

idx = np.array([[[0, 1, 1], [-1, -1, -1]],
                [[5, 1, 3], [1, -1, -1]]]) # n*h*w, where n = a*b,

这里-1表示不应用索引。我希望得到一个结果

res = np.array([[1+2+2, 0],
                [5+2+4, 2]]) # h*w

在实际实践中，我正在使用非常大的 3D 张量(n ~ 万亿)，并且 idx 非常稀疏(即大量 -1)。由于 idx 是固定的，我目前的解决方案是通过填充 0 和 1 来预先计算一个 n*(h*w) 数组 index_tensor，然后执行

tmp = arr.reshape(1, n)
res = (tmp @ index_tensor).reshape([h,w])

它工作正常，但需要大量内存来存储index_tensor。有没有什么方法可以利用 idx 的稀疏性和不可更改性来降低内存成本并在 python 中保持公平的运行速度(使用 numpy 或 pytorch 是最好的)？提前致谢!

Answer 1

忽略-1目前比较复杂，直接的索引和求和是:

In [58]: arr = np.array([[ 1, 2, 3, 4, 5], [5, 6, 7, 8, 9]])
In [59]: idx = np.array([[[0, 1, 1], [2, 4, 6]],
    ...:                 [[5, 1, 3], [1, -1, -1]]])
In [60]: arr.flat[idx]
Out[60]: 
array([[[1, 2, 2],
        [3, 5, 6]],

       [[5, 2, 4],
        [2, 9, 9]]])
In [61]: _.sum(axis=-1)
Out[61]: 
array([[ 5, 14],
       [11, 20]])

处理 -1 的一种方法(不一定快速或内存高效)是使用掩码数组:

In [62]: mask = idx<0
In [63]: mask
Out[63]: 
array([[[False, False, False],
        [False, False, False]],

       [[False, False, False],
        [False,  True,  True]]])

In [65]: ma = np.ma.masked_array(Out[60],mask)
In [67]: ma
Out[67]: 
masked_array(
  data=[[[1, 2, 2],
         [3, 5, 6]],

        [[5, 2, 4],
         [2, --, --]]],
  mask=[[[False, False, False],
         [False, False, False]],

        [[False, False, False],
         [False,  True,  True]]],
  fill_value=999999)
In [68]: ma.sum(axis=-1)
Out[68]: 
masked_array(
  data=[[5, 14],
        [11, 2]],
  mask=[[False, False],
        [False, False]],
  fill_value=999999)

掩码数组通过用中性值替换掩码值来处理求和之类的操作，例如求和时用 0 代替。

(我可能会在早上再看一遍)。

与矩阵乘积相加

In [72]: np.einsum('ijk,ijk->ij',Out[60],~mask)
Out[72]: 
array([[ 5, 14],
       [11,  2]])

这比掩码数组方法更直接、更快。

您还没有详细说明如何构建 index_tensor所以我不会尝试比较它。

另一种可能性是用 0 填充数组，并调整索引:

In [83]: arr1 = np.hstack((0,arr.ravel()))
In [84]: arr1
Out[84]: array([0, 1, 2, 3, 4, 5, 5, 6, 7, 8, 9])
In [85]: arr1[idx+1]
Out[85]: 
array([[[1, 2, 2],
        [3, 5, 6]],

       [[5, 2, 4],
        [2, 0, 0]]])
In [86]: arr1[idx+1].sum(axis=-1)
Out[86]: 
array([[ 5, 14],
       [11,  2]])

稀疏

首次尝试使用稀疏矩阵:

reshape idx到二维:

In [141]: idx1 = np.reshape(idx,(4,3))

从中创建一个稀疏张量。首先，我将迭代 lil方法，尽管通常构建 coo (甚至 csr )直接输入更快:

In [142]: M = sparse.lil_matrix((4,10),dtype=int)
     ...: for i in range(4):
     ...:     for j in range(3):
     ...:         v = idx1[i,j]
     ...:         if v>=0:
     ...:            M[i,v] = 1
     ...: 
In [143]: M
Out[143]: 
<4x10 sparse matrix of type '<class 'numpy.int64'>'
    with 9 stored elements in List of Lists format>
In [144]: M.A
Out[144]: 
array([[1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 1, 0, 1, 0, 1, 0, 0, 0],
       [0, 1, 0, 1, 0, 1, 0, 0, 0, 0],
       [0, 1, 0, 0, 0, 0, 0, 0, 0, 0]])

然后可以将其用于求乘积之和:

In [145]: <a href="https://stackoverflow.com/cdn-cgi/l/email-protection" class="__cf_email__" data-cfemail="e1aca1809393cf938097848d" rel="noreferrer noopener nofollow">[email protected]</a>()
Out[145]: array([ 3, 14, 11,  2])

使用<a href="https://stackoverflow.com/cdn-cgi/l/email-protection" class="__cf_email__" data-cfemail="cf82e18e8faebdbde1bdaeb9aaa3" rel="noreferrer noopener nofollow">[email protected]</a>()本质上就是你所做的。而M稀疏，arr不是。对于本例M.A@比 M@ 快.