给定一维索引数组，从 3D 数组向量化创建 2D 数组

Question

给定一个形状为 (2,n,m) 的数组 x 和一组索引 [i,j]，其中i,j ，我正在尝试获取形状为 (2,m) 的数组，其中第一个数组位于索引 [0,i]，第二个数组位于索引 [1,j]。这是一个用于泛化到形状为 (b,n,m) 的数组和一组长度为 b 的索引的测试用例。

此操作的明显选择是 np.choose，但这与预期的不同。我们希望将第一个数组的行 i 与第二个数组的行 j 配对。但是，当使用 np.choose([i,j],x) 时，np.choose 会与 中的第一个 column 配对具有索引 i 的数组 与具有索引 j 的 array 的第二列（可以在下面的代码中看到）到得到一个形状为 (n,m) 的数组。显然，使用 for 循环很容易执行此任务，但由于用例（在 Keras 中作为张量的自定义函数，禁止迭代），我不能这样做。是否有使用 Keras 后端函数或 Numpy 执行此操作的矢量化方式？我目前正在考虑使用“地图”来执行此操作，如果我弄明白了，我会用我自己的答案进行更新。

这里有一段代码 sn-p 向您展示 np.choose 如何处理 (2,n,m) 数组：

>>> import numpy as np 
>>> x = np.random.rand(2,4,2)
>>> choices = [3,1]
>>> np.choose(choices,x)
    ValueError: invalid entry in choice array
>>> np.choose([0,0],x)
    #Returns an array with x[0,:,0] and x[0,:,1] in shape(4,2)

Answer 1

我相信你可以使用advanced indexing。举个例子：

import numpy as np 
x = np.random.randint(0,10,(2,4,3))

x 是：

[[[0 4 1]
  [8 8 1]
  [3 3 6]
  [4 7 8]]

 [[7 1 2]
  [5 9 9]
  [0 4 0]
  [7 8 3]]]

现在x[[0,1],[3,1],:] 是：

[[4 7 8]
 [5 9 9]]

这可以扩展到(b,m,n) 问题：

import numpy as np 
x = np.random.randint(0,10,(100,200,300))
choices= np.random.randint(0,200,(100))

def loop():
    res=np.empty((100,300),int)
    for i in range(100):
        res[i]=x[i,choices[i]]
    return res    

还有一些性能测试：

In [30]: %timeit loop()
10000 loops, best of 3: 140 µs per loop

In [31]: %timeit x[arange(100),choices,:]
10000 loops, best of 3: 23.7 µs per loop

这里的索引方法只比循环快 6 倍，因为任务 （提取）不能利用内存对齐。

最后，您可以通过loop2=numba.njit(loop) 及时编译来增强循环。

In [32]: %timeit loop2()
10000 loops, best of 3: 32 µs per loop

这表明索引方法是最优的。