PyTorch：用二维张量索引高维张量答案

【问题标题】：PyTorch: Index high dimensional tensor with two dimensional tensorPyTorch：用二维张量索引高维张量
【发布时间】：2020-09-26 05:33:03
【问题描述】：

假设我有以下张量：

N = 2
k = 3
d = 2

L = torch.arange(N * k * d * d).view(N, k, d, d)
L
tensor([[[[ 0,  1],
          [ 2,  3]],

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

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


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

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

         [[20, 21],
          [22, 23]]]])


index = torch.Tensor([0,1,0,0]).view(N,-1)
index
tensor([[0., 1.],
        [0., 0.]])

我现在想使用索引张量在第二维上挑选出相应的矩阵，即我想得到类似的东西：

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

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


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

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

知道如何实现这一目标吗？非常感谢！

【问题讨论】：

标签： numpy indexing pytorch

【解决方案1】：

张量可以用跨不同维度指定的多个张量进行索引（张量的元组），其中每个张量的 i-th 元素组合起来创建一个索引元组，即data[indices_dim0, indices_dim1]结果在索引data[indices_dim0[0], indices_dim1[0]]、data[indices_dim0[1], indices_dim1[1]] 等。它们必须具有相同的长度len(indices_dim0) == len(indices_dim1)。

让我们使用index 的平面版本（在应用视图之前）。每个元素都需要与适当的批次索引匹配，即[0, 0, 1, 1]。 index 也需要有 torch.long 类型，因为浮点数不能用作索引。 torch.tensor 应该首选用于使用现有数据创建张量，因为 torch.Tensor 是默认张量类型 (torch.FloatTensor) 的别名，而 torch.tensor 自动使用表示给定值的数据类型，但也支持dtype 参数手动设置类型，一般比较通用。

# Type torch.long is inferred
index = torch.tensor([0, 1, 0, 0])

# Same, but explicitly setting the type
index = torch.tensor([0, 1, 0, 0], dtype=torch.long)

batch_index = torch.tensor([0, 0, 1, 1])

L[batch_index, index]
# => tensor([[[ 0,  1],
#             [ 2,  3]],
#
#            [[ 4,  5],
#             [ 6,  7]],
#
#            [[12, 13],
#             [14, 15]],
#
#            [[12, 13],
#             [14, 15]]])

索引不限于一维张量，但它们都需要具有相同的大小，并且每个元素都用作一个索引，例如对于二维张量，索引发生为data[indices_dim0[i][j], indices_dim1[i][j]] 使用 2D 张量，创建批次索引恰好要简单得多，而无需手动进行。

index = torch.tensor([0, 1, 0, 0]).view(N, -1)
# => tensor([[0, 1],
#            [0, 0]])

# Every batch gets its index and is repeated across dim=1
batch_index = torch.arange(N).view(N, 1).expand_as(index)
# => tensor([[0, 0],
#            [1, 1]])

L[batch_index, index]

【讨论】：