如何使用 torch.roll 移动张量的特定元素？

Question

我有一个张量 x，它看起来像这样：

x = tensor([  1,  2,  3,  4,  5],
           [  6,  7,  8,  9, 10]
           [ 11, 12, 13, 14, 15])

我正在尝试切换每个张量的前两个和最后两个数字，如下所示：

x = tensor([  4,  5,  3,  1,  2],
           [  9, 10,  8,  6,  7],
           [ 14, 15, 13, 11, 12])

如何使用 torch.roll() 执行此操作？我如何切换 3 而不是 1？

Answer 1

不确定单独使用 torch.roll 是否可以完成...但是，您可以通过使用临时张量和配对赋值来获得所需的结果：

>>> x = torch.arange(1, 16).reshape(3,-1)
tensor([[ 1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10],
        [11, 12, 13, 14, 15]])

>>> tmp = x.clone()

# swap the two sets of columns
>>> x[:,:2], x[:,-2:] = tmp[:,-2:], tmp[:,:2]

这样张量 x 已经变异为：

>>> x
tensor([[ 4,  5,  3,  1,  2],
        [ 9, 10,  8,  6,  7],
        [14, 15, 13, 11, 12]])

您可以使用 torch.roll 和一些索引来完成此操作：

>>> x = torch.arange(1, 21).reshape(4,-1)
tensor([[ 1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10],
        [11, 12, 13, 14, 15],
        [16, 17, 18, 19, 20]])

>>> rolled = x.roll(-2,0)
tensor([[11, 12, 13, 14, 15],
        [16, 17, 18, 19, 20],
        [ 1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10]])

# overwrite columns [1,-1[ from rolled with those from x
>>> rolled[:, 1:-1] = x[:, 1:-1]

在这一端你得到：

>>> rolled
tensor([[11,  2,  3,  4, 15],
        [16,  7,  8,  9, 20],
        [ 1, 12, 13, 14,  5],
        [ 6, 17, 18, 19, 10]])

如何使用 torch.roll 移动张量的特定元素？

How do I shift specific elements of a tensor with torch.roll?

python

pytorch

tensor