如何在不使用显式循环重组 Pytorch 中的张量行的情况下生成 [0,2,4,1,3,5] 之类的索引？

Question

假设我有一个像

这样的张量

a = torch.tensor([[3, 1, 5, 0, 4, 2],
                  [2, 1, 3, 4, 5, 0],
                  [0, 4, 5, 1, 2, 3],
                  [3, 1, 4, 5, 0, 2],
                  [3, 5, 4, 2, 0, 1],
                  [5, 3, 0, 4, 1, 2]])

我想通过应用变换 a[c] where

来重新组织张量的行

c = torch.tensor([0,2,4,1,3,5])

得到

b = torch.tensor([[3, 1, 5, 0, 4, 2],
                  [0, 4, 5, 1, 2, 3],
                  [3, 5, 4, 2, 0, 1],
                  [2, 1, 3, 4, 5, 0],
                  [3, 1, 4, 5, 0, 2],
                  [5, 3, 0, 4, 1, 2]])

为了做到这一点，我想生成张量 c 这样我就可以进行这个转换，而不管张量 a 的大小和步长大小（我在这里取等于 2为简单起见的示例）。谁能告诉我如何在 PyTorch 中不使用显式 for 循环的情况下为一般情况生成这样的张量？

Answer 1

想了一会儿，我想到了以下生成索引的解决方案

step = 2
idx = torch.arange(0,a.size(0),step)
# idx = tensor([0, 2, 4])
idx = idx.repeat(int(a.size(0)/idx.size(0)))
# idx = tensor([0, 2, 4, 0, 2, 4])
incr = torch.arange(0,step)
# incr = tensor([0, 1])
incr = incr.repeat_interleave(int(a.size(0)/incr.size(0)))
# incr = tensor([0, 0, 0, 1, 1, 1])
c = incr + idx
# c = tensor([0, 2, 4, 1, 3, 5])

这之后张量c可以通过使用

得到张量b

b = a[c.long()]

Answer 2

我也想出了另外一个方案，解决了上面重组张量a的行生成张量b而不生成索引数组c[=14]的问题=]

step = 2
b = a.view(-1,step,a.size(-1)).transpose(0,1).reshape(-1,a.size(-1))

Answer 3

可以使用torch.index_select，所以：

b = torch.index_select(a, 0, c)

官方文档解释的很清楚

如何在不使用显式循环重组 Pytorch 中的张量行的情况下生成 [0,2,4,1,3,5] 之类的索引？

How to generate indices like [0,2,4,1,3,5] without using explicit loop for reorganizing rows of a tensors in Pytorch?

python

numpy

pytorch

tensor