使用 2d 数组作为 4d 数组的索引

Question

我有一个来自 tensor.max() 操作的 Numpy 2D 数组 (4000,8000)，它存储 4D 数组 (30,4000,8000,3) 的第一个维度的索引。我需要获取一个 (4000,8000,3) 数组，该数组使用这组图像的索引并提取 2D max 数组中每个位置的像素。

A = np.random.randint( 0, 29, (4000,8000), dtype=int)
B = np.random.randint(0,255,(30,4000,8000,3),dtype=np.uint8)

final = np.zeros((B.shape[1],B.shape[2],3))

r = 0
c = 0
for row in A:
  c = 0
  for col in row:
    x = A[r,c]
    final[r,c] = B[x,r,c]
    c=c+1
  r=r+1

print(final.shape)

有没有矢量化的方式来做到这一点？我正在使用循环与 RAM 使用作斗争。谢谢

Answer 1

您可以使用np.take_along_axis。

首先让我们创建一些数据（您应该提供了reproducible example）：

>>> N, H, W, C = 10, 20, 30, 3
>>> arr = np.random.randn(N, H, W, C)
>>> indices = np.random.randint(0, N, size=(H, W))

然后，我们将使用np.take_along_axis。但为此，indices 数组的形状必须与arr 数组的形状相同。所以我们使用np.newaxis 在形状不匹配的地方插入轴。

>>> res = np.take_along_axis(arr, indices[np.newaxis, ..., np.newaxis], axis=0)

它已经提供了可用的输出，但在第一个轴上有一个单一维度：

>>> res.shape
(1, 20, 30, 3)

所以我们可以挤压它：

>>> res = np.squeeze(res)

>>> res.shape
(20, 30, 3)

并最终检查数据是否符合我们的要求：

>>> np.all(res[0, 0] == arr[indices[0, 0], 0, 0])
True

>>> np.all(res[5, 3] == arr[indices[5, 3], 5, 3])
True