PyTorch基础（14）-- torch.roll()方法

Posted 2021-12-11 奋斗丶

一、前言

在阅读swin transformer这篇论文时，论文中提出了一种较为新颖的创新点，shifted window技术，该技术可以解决窗口之间信息未交互的问题，有兴趣的可以去查阅swin transformer对CNN的降维打击。在阅读代码的过程中，发现shifted window技术的实现原理其实非常简单，仅仅用到了torch.roll()这个方法，相信你理解了该方法后，会进一步理解swin transformer的高明之处。话不多说，let’s go go go!

二、torch.roll()方法解析

从pytorch的官方文档中，我们可以找到torch.roll()方法的解释：沿着给定的维度滚动tensor。在第一个位置处，重新引入超过最后一个位置的元素（这句话非常拗口，大家后续结合实例，可以很清晰地明白这句话是什么意思，请看完下列实例之后再来理解）。如果没有指定一个维度，那么tensor首先被展开，然后重新恢复到原来的shape。此外，roll()方法有三个参数。

• input：输入的tensor
• shifts：可以为int，也可以是int型的元组。可以理解为roll的步数（参照大富翁游戏中的步数）。如果是tuple型，那么维度必须与tuple具有相同size，并按照每个维度roll相应的步数。
• dims：roll的维度

三、案例分析

3.1 例1 — shifts=1 & dims未指定

x = np.array([[1, 2, 3],
              [4, 5, 6],
              [7, 8, 9]])
x = torch.from_numpy(x)
print('before roll', x)
x = torch.roll(x, 1)
print('after roll', x)

3.2 例2 — shifts=1 & dims=0

x = np.array([[1, 2, 3],
              [4, 5, 6],
              [7, 8, 9]])
x = torch.from_numpy(x)
print('before roll', x)
x = torch.roll(x, 1, dims=0)
print('after roll', x)

3.3 例3 — shifts=1 & dims=1

x = np.array([[1, 2, 3],
              [4, 5, 6],
              [7, 8, 9]])
x = torch.from_numpy(x)
print('before roll', x)
x = torch.roll(x, 1, dims=1)
print('after roll', x)

3.4 例4 — shifts=(1,1) & dims=(0,1)

x = np.array([[1, 2, 3],
              [4, 5, 6],
              [7, 8, 9]])
x = torch.from_numpy(x)
print('before roll', x)
x = torch.roll(x, (1, 1), dims=(0,1))
print('after roll', x)

参考

torch.roll方法官方解释