pytorch - 在pytorch中使用索引提取张量数据

标签 pytorch data-manipulation tensor

我希望张量以某种方式索引。

假设我的数据，张量 X 形 (1, 3, 16, 9) 是

tensor([[[[ 0.,  0.,  0.,  0.,  1.,  2.,  0.,  5.,  6.],
      [ 0.,  0.,  0.,  1.,  2.,  3.,  5.,  6.,  7.],
      [ 0.,  0.,  0.,  2.,  3.,  4.,  6.,  7.,  8.],
      [ 0.,  0.,  0.,  3.,  4.,  0.,  7.,  8.,  0.],
      [ 0.,  1.,  2.,  0.,  5.,  6.,  0.,  9., 10.],
      [ 1.,  2.,  3.,  5.,  6.,  7.,  9., 10., 11.],
      [ 2.,  3.,  4.,  6.,  7.,  8., 10., 11., 12.],
      [ 3.,  4.,  0.,  7.,  8.,  0., 11., 12.,  0.],
      [ 0.,  5.,  6.,  0.,  9., 10.,  0., 13., 14.],
      [ 5.,  6.,  7.,  9., 10., 11., 13., 14., 15.],
      [ 6.,  7.,  8., 10., 11., 12., 14., 15., 16.],
      [ 7.,  8.,  0., 11., 12.,  0., 15., 16.,  0.],
      [ 0.,  9., 10.,  0., 13., 14.,  0.,  0.,  0.],
      [ 9., 10., 11., 13., 14., 15.,  0.,  0.,  0.],
      [10., 11., 12., 14., 15., 16.,  0.,  0.,  0.],
      [11., 12.,  0., 15., 16.,  0.,  0.,  0.,  0.]],

     [[ 0.,  0.,  0.,  0., 17., 18.,  0., 21., 22.],
      [ 0.,  0.,  0., 17., 18., 19., 21., 22., 23.],
      [ 0.,  0.,  0., 18., 19., 20., 22., 23., 24.],
      [ 0.,  0.,  0., 19., 20.,  0., 23., 24.,  0.],
      [ 0., 17., 18.,  0., 21., 22.,  0., 25., 26.],
      [17., 18., 19., 21., 22., 23., 25., 26., 27.],
      [18., 19., 20., 22., 23., 24., 26., 27., 28.],
      [19., 20.,  0., 23., 24.,  0., 27., 28.,  0.],
      [ 0., 21., 22.,  0., 25., 26.,  0., 29., 30.],
      [21., 22., 23., 25., 26., 27., 29., 30., 31.],
      [22., 23., 24., 26., 27., 28., 30., 31., 32.],
      [23., 24.,  0., 27., 28.,  0., 31., 32.,  0.],
      [ 0., 25., 26.,  0., 29., 30.,  0.,  0.,  0.],
      [25., 26., 27., 29., 30., 31.,  0.,  0.,  0.],
      [26., 27., 28., 30., 31., 32.,  0.,  0.,  0.],
      [27., 28.,  0., 31., 32.,  0.,  0.,  0.,  0.]],

     [[ 0.,  0.,  0.,  0., 33., 34.,  0., 37., 38.],
      [ 0.,  0.,  0., 33., 34., 35., 37., 38., 39.],
      [ 0.,  0.,  0., 34., 35., 36., 38., 39., 40.],
      [ 0.,  0.,  0., 35., 36.,  0., 39., 40.,  0.],
      [ 0., 33., 34.,  0., 37., 38.,  0., 41., 42.],
      [33., 34., 35., 37., 38., 39., 41., 42., 43.],
      [34., 35., 36., 38., 39., 40., 42., 43., 44.],
      [35., 36.,  0., 39., 40.,  0., 43., 44.,  0.],
      [ 0., 37., 38.,  0., 41., 42.,  0., 45., 46.],
      [37., 38., 39., 41., 42., 43., 45., 46., 47.],
      [38., 39., 40., 42., 43., 44., 46., 47., 48.],
      [39., 40.,  0., 43., 44.,  0., 47., 48.,  0.],
      [ 0., 41., 42.,  0., 45., 46.,  0.,  0.,  0.],
      [41., 42., 43., 45., 46., 47.,  0.,  0.,  0.],
      [42., 43., 44., 46., 47., 48.,  0.,  0.,  0.],
      [43., 44.,  0., 47., 48.,  0.,  0.,  0.,  0.]]]]

我想要那些 (row_index % n) == i 的行(比如 n = 4 和 i = 0 到 3) 保存在另一个张量 Y 中。

例如，对于数据X[0][0]:

[[ 0.,  0.,  0.,  0.,  1.,  2.,  0.,  5.,  6.],
 [ 0.,  0.,  0.,  1.,  2.,  3.,  5.,  6.,  7.],
 [ 0.,  0.,  0.,  2.,  3.,  4.,  6.,  7.,  8.],
 [ 0.,  0.,  0.,  3.,  4.,  0.,  7.,  8.,  0.],
 [ 0.,  1.,  2.,  0.,  5.,  6.,  0.,  9., 10.],
 [ 1.,  2.,  3.,  5.,  6.,  7.,  9., 10., 11.],
 [ 2.,  3.,  4.,  6.,  7.,  8., 10., 11., 12.],
 [ 3.,  4.,  0.,  7.,  8.,  0., 11., 12.,  0.],
 [ 0.,  5.,  6.,  0.,  9., 10.,  0., 13., 14.],
 [ 5.,  6.,  7.,  9., 10., 11., 13., 14., 15.],
 [ 6.,  7.,  8., 10., 11., 12., 14., 15., 16.],
 [ 7.,  8.,  0., 11., 12.,  0., 15., 16.,  0.],
 [ 0.,  9., 10.,  0., 13., 14.,  0.,  0.,  0.],
 [ 9., 10., 11., 13., 14., 15.,  0.,  0.,  0.],
 [10., 11., 12., 14., 15., 16.,  0.,  0.,  0.],      
 [11., 12.,  0., 15., 16.,  0.,  0.,  0.,  0.]]

我想要一个包含以下数据的张量，它基本上是其中 row_index % 4 == 0 的行的集合(此处 i = 0): 

[[ 0.,  0.,  0.,  0.,  1.,  2.,  0.,  5.,  6.],
 [ 0.,  1.,  2.,  0.,  5.,  6.,  0.,  9., 10.],
 [ 0.,  5.,  6.,  0.,  9., 10.,  0., 13., 14.],
 [ 0.,  9., 10.,  0., 13., 14.,  0.,  0.,  0.]]

类似地，其中 i = 1、row_index % 4 == i 将如下所示:

[[ 0.,  0.,  0.,  1.,  2.,  3.,  5.,  6.,  7.],
 [ 1.,  2.,  3.,  5.,  6.,  7.,  9., 10., 11.],
 [ 5.,  6.,  7.,  9., 10., 11., 13., 14., 15.],
 [ 9., 10., 11., 13., 14., 15.,  0.,  0.,  0.]]

当i = 2时，row_index % 4 == i:

[[ 0.,  0.,  0.,  2.,  3.,  4.,  6.,  7.,  8.],
 [ 2.,  3.,  4.,  6.,  7.,  8., 10., 11., 12.],
 [ 6.,  7.,  8., 10., 11., 12., 14., 15., 16.],
 [10., 11., 12., 14., 15., 16.,  0.,  0.,  0.]]

当i = 3时，row_index % 4 == i:

[[ 0.,  0.,  0.,  3.,  4.,  0.,  7.,  8.,  0.],
 [ 3.,  4.,  0.,  7.,  8.,  0., 11., 12.,  0.],
 [ 7.,  8.,  0., 11., 12.,  0., 15., 16.,  0.],
 [11., 12.,  0., 15., 16.,  0.,  0.,  0.,  0.]]

我尝试过对其进行硬编码，但当数据变大并且大小变得动态时，它似乎不切实际，我认为会有更好的方法来实现它。

temp0 = data[0][0][0][:] 
temp1 = data[0][0][4][:]
temp2 = data[0][0][8][:]
temp3 = data[0][0][12][:]
temp = torch.stack([temp0,temp1,temp2,temp3],dim = 0)

此外，如果结果能够以一个张量形式返回，那就太好了:

tensor Y = ([[[ 0.,  0.,  0.,  0.,  1.,  2.,  0.,  5.,  6.],
              [ 0.,  1.,  2.,  0.,  5.,  6.,  0.,  9., 10.],
              [ 0.,  5.,  6.,  0.,  9., 10.,  0., 13., 14.],
              [ 0.,  9., 10.,  0., 13., 14.,  0.,  0.,  0.]], 

             [[ 0.,  0.,  0.,  1.,  2.,  3.,  5.,  6.,  7.],
              [ 1.,  2.,  3.,  5.,  6.,  7.,  9., 10., 11.],
              [ 5.,  6.,  7.,  9., 10., 11., 13., 14., 15.],
              [ 9., 10., 11., 13., 14., 15.,  0.,  0.,  0.]], 
   
             [[ 0.,  0.,  0.,  2.,  3.,  4.,  6.,  7.,  8.],
              [ 2.,  3.,  4.,  6.,  7.,  8., 10., 11., 12.],
              [ 6.,  7.,  8., 10., 11., 12., 14., 15., 16.],
              [10., 11., 12., 14., 15., 16.,  0.,  0.,  0.]], 

             [[ 0.,  0.,  0.,  3.,  4.,  0.,  7.,  8.,  0.],
              [ 3.,  4.,  0.,  7.,  8.,  0., 11., 12.,  0.],
              [ 7.,  8.,  0., 11., 12.,  0., 15., 16.,  0.],
              [11., 12.,  0., 15., 16.,  0.,  0.,  0.,  0.]]])

最佳答案

您可以通过首先构造包含所选行的张量，然后使用 torch.gather 来实现此目的组装最终张量。

假设我们有两个列表I和N，其中包含i和n<的值 分别:

I = [0, 1, 2, 3]
N = [4, 4, 4, 4]

首先我们构造索引张量:

>>> index = torch.stack([(torch.arange(16) % n == i).nonzero() for i, n in zip(I, N)])
tensor([[[ 0],
         [ 4],
         [ 8],
         [12]],

        [[ 1],
         [ 5],
         [ 9],
         [13]],

        [[ 2],
         [ 6],
         [10],
         [14]],

        [[ 3],
         [ 7],
         [11],
         [15]]])

然后需要进行一些扩展和 reshape :

>>> index_ = index[None].flatten(1,2).expand(X.size(0), -1, X.size(-1))
tensor([[[ 0,  0,  0,  0,  0,  0,  0,  0,  0],
         [ 4,  4,  4,  4,  4,  4,  4,  4,  4],
         [ 8,  8,  8,  8,  8,  8,  8,  8,  8],
         [12, 12, 12, 12, 12, 12, 12, 12, 12],
         [ 1,  1,  1,  1,  1,  1,  1,  1,  1],
         [ 5,  5,  5,  5,  5,  5,  5,  5,  5],
         [ 9,  9,  9,  9,  9,  9,  9,  9,  9],
         [13, 13, 13, 13, 13, 13, 13, 13, 13],
         [ 2,  2,  2,  2,  2,  2,  2,  2,  2],
         [ 6,  6,  6,  6,  6,  6,  6,  6,  6],
         [10, 10, 10, 10, 10, 10, 10, 10, 10],
         [14, 14, 14, 14, 14, 14, 14, 14, 14],
         [ 3,  3,  3,  3,  3,  3,  3,  3,  3],
         [ 7,  7,  7,  7,  7,  7,  7,  7,  7],
         [11, 11, 11, 11, 11, 11, 11, 11, 11],
         [15, 15, 15, 15, 15, 15, 15, 15, 15]]])

根据经验，我们希望 index_ 具有与 X 相同的维度。

现在我们可以应用torch.gather并 reshape 为最终形式:

>>> X.gather(1, index_).reshape(len(X), *index.shape[:2], -1)
tensor([[[[ 0.,  0.,  0.,  0.,  1.,  2.,  0.,  5.,  6.],
          [ 0.,  1.,  2.,  0.,  5.,  6.,  0.,  9., 10.],
          [ 0.,  5.,  6.,  0.,  9., 10.,  0., 13., 14.],
          [ 0.,  9., 10.,  0., 13., 14.,  0.,  0.,  0.]],

         [[ 0.,  0.,  0.,  1.,  2.,  3.,  5.,  6.,  7.],
          [ 1.,  2.,  3.,  5.,  6.,  7.,  9., 10., 11.],
          [ 5.,  6.,  7.,  9., 10., 11., 13., 14., 15.],
          [ 9., 10., 11., 13., 14., 15.,  0.,  0.,  0.]],

         [[ 0.,  0.,  0.,  2.,  3.,  4.,  6.,  7.,  8.],
          [ 2.,  3.,  4.,  6.,  7.,  8., 10., 11., 12.],
          [ 6.,  7.,  8., 10., 11., 12., 14., 15., 16.],
          [10., 11., 12., 14., 15., 16.,  0.,  0.,  0.]],

         [[ 0.,  0.,  0.,  3.,  4.,  0.,  7.,  8.,  0.],
          [ 3.,  4.,  0.,  7.,  8.,  0., 11., 12.,  0.],
          [ 7.,  8.,  0., 11., 12.,  0., 15., 16.,  0.],
          [11., 12.,  0., 15., 16.,  0.,  0.,  0.,  0.]]]])

该方法可以扩展到批量张量:

>>> index = torch.stack([(torch.arange(16) % n == i).nonzero() for i, n in zip(I, N)])
>>> index_  = index[None,None].flatten(2,3).expand(X.size(0), X.size(1), -1, X.size(-1))

>>> X.gather(2, index_).reshape(*X.shape[:2], *index.shape[:2], -1)

关于pytorch - 在pytorch中使用索引提取张量数据，我们在Stack Overflow上找到一个类似的问题： https://stackoverflow.com/questions/69938529/

上一篇：r - 绘制跨越180度国际日期变更线的凸包并计算面积

下一篇：android - 从类(class)打开 Activity

相关文章：

image - 在同一图像张量上两次使用 torchvision.utils.save_image 会使第二次保存无效。这是怎么回事？

pytorch - 从掩码索引获取张量列表

python - 为什么PIL与Pytorch经常使用？

pytorch - 如何在 PyTorch 中的特定新维度中重复张量

用其他列的值替换列中的 NA

pytorch - 如何用 2D mask mask 3D 张量并保持原始向量的尺寸？

python - 为什么有些人将两个不同网络的参数链接起来并使用相同的优化器进行训练？

python - 在 PyTorch 中将张量的一部分分割成另一个张量的最快方法是什么？

python - Pandas - 将多个组行合并为一行

python - Pandas 中的数据处理

©2024 IT工具网  联系我们