Numpy PyTorch：具有二维张量的索引高维张量

假设我有以下张量：

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]]])

你知道我怎样才能做到这一点吗？

---

非常感谢你

可以使用跨不同维度的张量元组指定的多个张量对张量进行索引，其中每个张量的第i个元素被组合以创建一个索引元组，即数据[Indexes_dim0，Indexes_dim1]导致索引数据[Indexes_dim0[0]，Indexes_dim1[0]]，数据[Indexes_dim0[1]，Indexes_dim1[1]等等。它们必须具有相同的长度lenindices_dim0==lenindices_dim1

在应用视图之前，让我们使用索引的平面版本。每个元素都需要匹配到适当的批处理索引，即[0,0,1,1]。此外，索引需要具有torch.long类型，因为浮点数不能用作索引。应首选使用现有数据创建张量，因为torch.Tensor是默认张量类型torch.FloatTensor的别名，而torch.Tensor自动使用表示给定值的数据类型，但也支持dtype参数手动设置类型，通常更通用

推断出torch.long类型 指数=火炬张量[0,1,0,0] 相同，但显式设置类型 index=torch.tensor[0,1,0,0]，dtype=torch.long 批次指数=火炬张量[0,0,1,1] L[批次索引，索引] =>张量[[0,1]， [ 2, 3]], [[ 4, 5], [ 6, 7]], [[12, 13], [14, 15]], [[12, 13], [14, 15]]] 索引不限于一维张量，但它们都需要具有相同的大小，并且每个元素都用作一个索引，例如，对于二维张量，索引发生在数据[index_dim0[i][j]，index_dim1[i][j]] 使用2D张量，创建批次索引而无需手动操作，恰好要简单得多

index=torch.tensor[0,1,0,0].viewN，-1 =>张量[[0,1]， [0, 0]] 每个批次都获得其索引，并在dim=1中重复 批次索引=torch.arangeN.viewN，1.expand\u asindex =>张量[[0,0]， [1, 1]] L[批次索引，索引]