PyTorch：如何使用多维索引从多维张量中选择值？ [复制]

Question

给定一个形状为 (5, 2, 2) 的张量 input：

tensor([[[ 0,  1],
         [ 2,  3]],

        [[ 4,  5],
         [ 6,  7]],

        [[ 8,  9],
         [10, 11]],

        [[12, 13],
         [14, 15]],

        [[16, 17],
         [18, 19]]])

还有一个形状为 (2,2) 的张量 index：

tensor([[4, 3],
        [4, 2]])

如何获得以下输出：

tensor([[16, 13],
        [18, 11]])

具体来说：我有 5 个大小为 2x2 像素的输入图像（存储在 input 中）。现在我想将这 5 个图像组合成一个 2x2 像素大小的图像，其中index 确定每个像素应该从哪个输入图像复制它。

示例：从左上角index[0,0] == 4 开始，我取像素值input[4,0,0] == 16。然后继续index[0,1] == 3，我取像素值input[3,0,1] == 13等等。

Answer 1

我找到的解决方法和this one差不多：

indices = np.indices(input.shape)
indices[0] = index
input[tuple(indices)][0]

给出所需的输出

tensor([[16, 13],
        [18, 11]])