torch.roll() 用法解读

 torch.roll(input, shifts, dims=None) -> Tensor

定义

沿给定维度滚动张量输入。超出最后一个位置的元素在第一个位置重新引入。如果 dims 为 None，则张量将在滚动前被展平，然后恢复到原始形状。

参数

• input (Tensor) - 输入张量
• shifts (int or tuple of python: ints) - 张量元素移动的位置数。如果 shifts 是元组，则 dims 必须是大小相同的元组，每个维度都会滚动对应的值
• dims (int or tuple of python: ints) - 沿着dims指定的维度进行滚动

示例

import torch

x = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]).view(4, 2)
print(x)
# tensor([[1, 2],
#         [3, 4],
#         [5, 6],
#         [7, 8]])
x1 = torch.roll(x, 1)
print(x1)
# tensor([[8, 1],
#         [2, 3],
#         [4, 5],
#         [6, 7]])
x2 = torch.roll(x, 1, 0)
print(x2)
# tensor([[7, 8],
#         [1, 2],
#         [3, 4],
#         [5, 6]])
x3 = torch.roll(x, -1, 0)
print(x3)
# tensor([[3, 4],
#         [5, 6],
#         [7, 8],
#         [1, 2]])
x4 = torch.roll(x, shifts=(2, 1), dims=(0, 1))
print(x4)
# tensor([[6, 5],
#         [8, 7],
#         [2, 1],
#         [4, 3]])

参考

torch.roll — PyTorch 1.12 documentation