大张量上的 numpy n 维智能索引 - 内存效率

Question

我正在处理大型张量，因此临时张量的 numpy 内存分配开始显着影响执行时间 + 代码有时会在这些中间步骤中引发内存分配错误。这是我想出的用另一个张量的 int 值（例如 result_ijk = a[i, b[i, j], k]）索引一个张量的两种方法，尽管第二种方法看起来更节省内存，但我还是想创建这个巨大的索引矩阵和遍历它的所有值（甚至并行）有点有线（并且经常达到内存限制）：

def test():
    i, j, k, l = 10, 20, 30, 40 # in reality, they're like 1e3..1e6
    a = np.random.rand(i, j, k)
    b = np.random.randint(0, j, size=i*l).reshape((i, l))
    # c_ilk = c[i, b[i, l], k]; shape(c) = (10, 40, 30)
    tmp = a[:, b, :] # <= i*ijk additional memory allocated (!) crazy
    c1 = np.diagonal(tmp, axis1=0, axis2=1).transpose([2, 0, 1])
    print(c1.shape)
    # another approach:
    ii, ll = np.indices((i, l)) # <= 2*i*l of temporary ints allocated
    tmp2 = b[ii, ll] # i*l of ints allocated, slow ops
    c2 = a[ii, tmp2] # slow ops over tensor
    print(c2.shape)
    print(np.allclose(c1, c2))

test()

- 关于如何优化这种类型的 n-dim 智能索引代码有什么建议吗？

如果我要在 Theano 中使用这段 ~vectorized 代码，它是否也会分配所有这些临时缓冲区或者它可以设法构建它们 "on-fly"？是否有任何包可以以 lazy\more 有效的方式执行此类索引，而无需分配这些 ii 之类的张量？

（注意：我最后需要对其进行渐变，所以我不能使用像 numba 这样花哨的 jit 编译器 :( )

Answer 1

一些方法：

i,j,k,l=[100]*4
a = np.random.randint(0,5,(i, j, k))
b = np.random.randint(0, j,(i, l))

def test1():
    # c_ilk = c[i, b[i, l], k]; shape(c) = (2,3,5)
    tmp = a[:, b, :] # <= i*ijk additional memory allocated (!) crazy
    c1 = np.diagonal(tmp, axis1=0, axis2=1).transpose([2, 0, 1])
    return c1

def test2():
    ii, ll = np.indices((i, l)) # <= 2*i*l of temporary ints allocated
    tmp2 = b[ii, ll] # i*l of ints allocated, slow ops
    c2 = a[ii, tmp2] # slow ops over tensor
    #print(c2.shape)
    return c2

def test3():
    c3=np.empty((i,l,k),dtype=a.dtype)   
    for ii in range(i):
        for ll in range(l):
                c3[ii,ll]=a[ii,b[ii,ll]]
    return c3        

from numba import jit
test4=jit(test3)

以及相应的基准：

In [54]: %timeit test1()
1 loop, best of 3: 720 ms per loop

In [55]: %timeit test2()
100 loops, best of 3: 7.79 ms per loop

In [56]: %timeit test3()
10 loops, best of 3: 43.7 ms per loop

In [57]: %timeit test4()
100 loop, best of 3: 4.99 ms per loop

这似乎表明（请参阅@Eelco Hoogendoorn 评论）你的第二种方法几乎是大尺寸的最佳选择，而第一种是一个糟糕的选择。

对于 numba，您可以只使用这部分代码，并在非 "jited" 函数中应用渐变。

Answer 2

您只需分配一个长度为 i 的整数数组即可获得您想要的结果：

i_idx = np.arange(i)
c = a[i_idx[:, None], b[i_idx, :], :]
# or you can use the terser c = a[i_idx[:, None], b[i_idx]]

广播负责根据需要即时复制值，而无需为它们分配内存。

如果你为大型数组计时，你会发现它只比你的第二种方法快一点点：正如其他人所指出的，中间索引数组将比你的整体小几个数量级计算，因此对其进行优化对总运行时间或内存占用量影响很小。

大张量上的 numpy n 维智能索引 - 内存效率

numpy n-dimensional smart indexing over large tensors - memory efficiency

indexing

numpy

theano