PyTorch之PixelShuffle

这里介绍一下PyTorch中torch.nn.PixelShuffle()这个层的背景介绍和相关用法。

参考文档：
①PyTorch中的PixelShuffle

1 背景介绍

PixelShuffle层又名亚像素卷积层，是论文Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network中介绍的一种应用于超分辨率重建应用的具有上采样功能的卷积层。这篇ESPCN论文介绍了这种层的功能，sub-pixel convolution layer以$stride=\frac{1}{r}$($r$为SR的放大倍数upscaling factor)去提取feature map，虽然称之为卷积，但是其并没用用到任何需要学习的参数，它的原理也很简单，就是将输入feature map进行像素重组，也就是说亚像素卷积层虽用卷积之名，但却没有做任何乘加计算，只是用了另一种方式去提取特征罢了：

如上图所示的最后一层就是亚像素卷积层，它就是将输入格式为$(batch,r^{2}C,H,W)$的feature map中同一通道的像素提取出来作为输出feature map的一小块，遍历整个输入feature map就可以得到最后的输出图像。整体来看，就好像是用$\frac{1}{r}$的步长去做卷积一样，这样就造成了不是对整像素点做卷积，而是对亚像素做卷积，故称之为亚像素卷积层，最后的输出格式就是$(batch,1,rH,rW)$。
Note：

1. 想要了解更多关于ESPCN网络或者亚像素卷积层的相关知识，请看我的另一篇论文笔记之ESPCN。
2. 亚像素卷积是一个隐式卷积的过程，其不含任何可学习参数。

因此，简单一句话，PixelShuffle层做的事情就是将输入feature map像素重组输出高分辨率的feature map，是一种上采样方法，具体表达为：
$$(batch,r^{2}C,H,W)\to (batch,1,rH,rW)$$其中$r$为上采样倍率(上图中$r=3$)。

2 用法简介

我们来看看PyTorch是如何将ESPCN论文的核心——亚像素采样层编写出来的：
调用格式，该层的输入就是上采样倍率$r$：

torch.nn.PixelShuffle(upscale_factor=3)

内部源码如下：

其源码中就是3个方法：
①：init()：保存上采样倍率值。
②：forward()，调用F.pixel_shuffle()函数，其实做的就是亚像素卷积层背后的重组过程，其输入就是待重组的输入feature map。
③：extra_repr()：打印出上采样倍率$r$的值。

Note：

1. torch.nn.PixelShuffle()不含任何可学习参数$W,b$，和一般的卷积层是不一样的，尽管亚像素卷积也是一种特征提取过程。

2.1 实战代码

接下来我们实际来使用一下：
我们设置上采样倍率$r=2$，输入feature map的格式为：$(1,1,3,3)$，模拟一个batch=1，通道数为1，图像高和宽为3的feature map：

r = 2  # 上采样倍率
ps = nn.PixelShuffle(r)
x = torch.arange(4*9).view(1, 1*(r**2), 3, 3)
print(f'input is \n {x},and size is \n {x.size()}')
y = ps(x)  # 亚像素采样
print(f'output is \n {y},and size is \n {y.size()}')
print(f'upscale_factor is {ps.extra_repr()}')

2.2 效果展示

最后的输出如下：

可以很容易看出来，torch.nn.PixelShuffle()描述的就是ESPCN论文中的亚像素采样过程。