【Pytorch小知识】torch.roll()函数的用法及在Swin Transformer中的应用(详细易懂)

广告位：

图像拼接论文精读专栏 —— 图像拼接领域论文全覆盖（包含数据集），省时省力读论文，带你理解晦涩难懂的论文算法，学习零散的知识和数学原理，并学会写图像拼接领域的论文（介绍、相关工作、算法、实验、结论、并附有参考文献，不用一篇一篇文章再找）

图像拼接论文源码精读专栏 —— 图像拼接有源码的论文全覆盖（有的自己复现），帮助你通过源码进一步理解论文算法，助你做实验，跑出拼接结果，得到评价指标RMSE、SSIM、PSNR等，并寻找潜在创新点和改进提升思路。

超分辨率重建专栏 —— 从SRCNN开始，带你读论文，写代码，复现结果，找创新点，完成论文。手把手教，保姆级攻略。帮助你顺利毕业，熟练掌握超分技术。

图像去噪专栏 —— 从DnCNN开始，100个经典的基于深度学习的图像去噪算法。读懂论文，看懂代码，复现结果，寻找创新，新手小白入门，保姆级攻略。帮助你顺利毕业，实现目标。

关注文章最底部微信公众号【十小大的底层视觉工坊】，获取最精炼版论文解读！

有需要的同学可以点上面链接看看。

torch.roll()的官方文档

函数形式及参数说明

函数形式：torch.roll(input,shifts,dims=None）
参数解释：

• input:输入张量
• shifts:滚动的方向和长度，若为正，则向下滚动；若为负，则向上滚动。可是一个整数也可以是一个元组。(具体见下面例子)
• dims:张量按什么维度滚动。以二维为例，0就是按行上下或者下上(shifts为负)滚动；1就是左右或者右左(shifts为负)。可以是整数或元组。

举个例子

下面我们通过官方给出的例子加以讲解。

• x是一个定义好的4×2张量
• torch.roll(x,1,0):shifts=1，代表正向(向下)滚动1格，dims=0代表按行滚动。于是，[1,2]就到了第二行，[7,8]就到了第一行。以此类推。
• torch.roll(x,-1,0):shifs=-1，代表负向(向上)滚动1个，dims=0代表按行滚动。于是就有了与上一步相反的结果，每行都同步上移了一格。
• torch.roll(x,shifts=(2,1),dims=(0,1)):元组形式传参，那就是先按行正向滚动2格，再按列正向(左右)滚动1格。见如下过程：

x是这样的4×2张量:
$$\left[\begin{array}{cc}1& 2\\ 3& 4\\ 5& 6\\ 7& 8\end{array}\right]$$
shifts=2，dims=0后:
$$\left[\begin{array}{cc}5& 6\\ 7& 8\\ 1& 2\\ 3& 4\end{array}\right]$$
shifts=1，dims=1后：
$$\left[\begin{array}{cc}6& 5\\ 8& 7\\ 2& 1\\ 4& 3\end{array}\right]$$

torch.roll()的应用

Swin Transformer中的shifted window partition操作中更有效地计算自注意力的方法。在实现cyclic shift时要用到torch.roll()。

具体步骤

假设原图window partition是个4×4tensor：

接下来要实现上图中所示的A、B、C变换，即1到右下角；5，9，13到1的上面；2，3，4到1的左边；其他元素相对位置不变。那么可以通过以下两步实现：


这样就完成了cyclic shift。

若要完成reverse cyclic shift，则将对应的shifts改成正数即可实现。

没有硬件条件，需要云服务的同学可以扫码看看：