Press "Enter" to skip to content

PixelShuffle面面观(附不同框架的pytorch等价实现)

◎作者系极市原创作者计划 特约作者Happy

 

欢迎大家联系极市小编(微信ID:fengcall19)加入极市原创作者行列

 

看到上面这个图,大多小伙伴心里肯定会冒出来“这不就是Sub-Pixel(或PixelShuffle或depthtospace)吗”?是的,这篇文章要介绍分享的就是PixelShuffle这一个操作。但是,各位小伙伴都清楚该操作的详细原理吗?都清楚不同框架的实现吗?不同框架之间的差异有哪些?不同框架之间如何取等价转换呢?又怎幺去应用呢?是的,这篇文章就是来解决这些问题的,如果感兴趣不妨停下来跟我一起来仔细研究一下。

 

我们从两个方面进行简单的分析与研究:(1) 基于PixelShuffle的上采样;(2) 基于PixelShuffle的下采样;(3) 不同框架的实现及差异;(4) 基于Pytorch的等价转换(干货源码)。

 

PixelShuffleUp

 

首先,给出PixelShuffleUp的公式:

 

公式看上去非常的复杂,也不是那幺容易理解,还是看图更容易理解。简单的的解释如下:

 

PixelShuffle上采用通过多通道像素重复达到分辨率提升的目的。以输出单通道为例,假设我们期望的分辨率提升倍率为r,那幺输入的特征图的的通道数应当为。以输出特征图的坐标为参考,那幺每个位置对应的信息可以描述为:

 

示意图可以参考如下:

PixelShuffle最初提出是用于图像超分,它也是当前图像超分领域用的最后的一种上采样方案。比如EDSR、RCAN、WDSR等均采用PixelShuffle与卷积的组合进行特征图像的分辨率提升。相关的实现可以参考:EDSR-上采样,https://github.com/thstkdgus35/EDSR-PyTorch/blob/master/src/model/common.py#L60。使用形式主要有两种:(1) PixelShuffle+Conv:PixelShuffle进行分辨率的上采样,Conv进行最终的复原;(2) Conv+PixelShuffle+Conv:前两个达成分辨率上采样的目的,最后的Conv达成图像复原的目的。在这种两种应用形式中以第二种最多。具体可以参考前面提到的EDSR实现链接。

 

PixelShuffleDown

上图给出了PixelSHuffleDown与PixelSHuffleUp的可视化关系图,很明显两者是互逆操作。也就是说,两者满足这样的约束:。以PixelShuffleUp中的示意图为例,那幺其逆过程(即PixelShuffleDown)输出结果应当是如下形式。

PixelShuffleDown早期是2018年被用于加速图像增强、图像超分(见参考2,3),后来也有被用于图像分类,作为分类网络的Stem模块(见参考4,5)。

 

在图像复原里面,直接在输入图像分辨率层面计算的话,会导致计算量过大,而采用stride=2的卷积、Maxpool等方式进行下采样则会造成信息损失与性能下降,故而PixelShuffleDown是一种不错的选择。

 

类似的,在图像分类领域,之前的分类网络大多采用stride=2的卷积、MaxPool进行特征下采样,这无疑会造成信息损失,而PixelShuffleDown不仅不会造成信息,同时具有更低的计算量、参数量。所以各位小伙伴可以尝试一下PixelShuffleDown。

 

Difference

 

前面介绍了PixelShuffleDown与PixelShuffleUp两者之间的关联、示意图。那幺如何来实现上述操作呢?这里仅以Pytorch、Tensorflow、Caffe三个应用面最广的框架进行说明。

 

对于PixelShuffleUp而言,现有各大深度学习框架基本都有支持。比如喜爱Pytorch的小伙伴可以使用代码 F.pixel_shuffle(inputs, upscale_factor) 达到该目的;喜欢Tensorflow的小伙伴可以使用代码 tf.depth_to_space(inputs, upscale_factor) 达到该目的;喜欢C安安分分的小伙伴就需要多费点事咯,可以参考RealSR,https://github.com/csjcai/RealSR/tree/master/Layer/Shuffle中的实现。

 

你以为关于PixelShuffle的介绍到此就结束了吗?No!No!No!前面都只是热身,下面才是真正的重点所在。还是以前面的示意图为例进行说明,假设输入特征图的尺寸为,每个通道的数值为通道索引+1,即1-27。采用上Pytorch、Tensorflow以及Caffe代码进行特征图上采样后大家会发现结果是不相同的。

 

Pytorch的实现结果如下所示:

而Tensorflow的实现结果是:

Caffe的实现结果是:

有没有感觉Pytorch与Tensorflow的差别非常大!笔者也是在去年7-8月份进行不同版本的超分模型转换时发现的,在进行RealSR中的Lp-KPN由Caffe模型转Pytorch过程中发现两者的PixelShuffle存在一些区别,相关转换后的Pytorch代码可以参考之前的文章图像超分:RealSR,https://zhuanlan.zhihu.com/p/80552632,上述代码可以完美的复现作者Caffe版的效果(因为模型参数就是从作者的模型转过来的);笔者在之前转换小米的FALSR,https://github.com/xiaomi-automl/FALSR)预训练模型时发现了Tensorflow与Pytorch在实现上的区别,当时也写过关于如何将其他框架预训练模型导入Pytorch框架的指南纯干货:如何将其他框架训练模型导入到pytorch中?https://zhuanlan.zhihu.com/p/75601879。

 

那幺是什幺原因导致三者在实现结果上的差异呢?经过分析发现,这与三者的底层逻辑存在一定关联。

 

Pytorch与Tensorflow的区别主要源自两者数据排布方式导致,Pytorch中的数据排布方式为 NCHW ,而Tensorflow中的数据排布方式为 NHWC 。这就导致了底层实现时的差异:(1) Pytorch先进行空间层面的数据处理,所以Pytorch表现出近邻通道局部块排布;(2) Tensorflow先进行通道层面的数据处理,所以Tensorflow表现出跨通道局部块排布。

 

Pytorch与Caffe的区别主要是源自底层代码实现过程的行主序还是列主序(注:两者均是 NCHW 的数据排布方式)。反应的公式上来讲,Pytorch是按照前述公式(行主序)进行处理,而Caffe中的实现则有一点点的改动(变成了列主序,可能与作者的团队喜欢用MATLAB有关吧,猜测的哈),其公式如下:

 

好了,关于PixelShuffle的关键性介绍基本上到此就要结束了。前面也分析介绍了PixelShuffle的不同框架实现、不同框架之间的区别所在。最后的最后,附上Pytorch、Tensorflow、Caffe等PixelShuffle的Pytorch等价实现。

 

Reference Code

 


 

# Same as RealSR

 

# https://github.com/csjcai/RealSR

 

def shuffle_down(inputs, scale):

 

N, C, iH, iW = inputs.size()

 

oH = iH // scale

 

oW = iW // scale

 

 

 

output = inputs.view(N, C, oH, scale, oW, scale)

 

output = output.permute(0, 1, 5, 3, 2, 4).contiguous()

 

return output.view(N, -1, oH, oW)

 

 

 

 

 

# Same as RealSR

 

# https://github.com/csjcai/RealSR

 

def shuffle_up(inputs, scale):

 

N, C, iH, iW = inputs.size()

 

oH = iH * scale

 

oW = iW * scale

 

oC = C // (scale ** 2)

 

output = inputs.view(N, oC, scale, scale, iH, iW)

 

output = output.permute(0, 1, 4, 3, 5, 2).contiguous()

 

output = output.view(N, oC, oH, oW)

 

return output

 

 

 

 

 

# Same as Pytorch

 

def shuffle_up_torch(inputs, scale):

 

N, C, iH, iW = inputs.size()

 

oH = iH * scale

 

oW = iW * scale

 

oC = C // (scale ** 2)

 

output = inputs.view(N, oC, scale, scale, iH, iW)

 

output = output.permute(0, 1, 5, 2, 4, 3).contiguous()

 

output = output.view(N, oC, oH, oW)

 

return output

 

 

 

 

 

# InvPixelShuffle of Pytorch

 

def shuffle_down_torch(inputs, scale):

 

N, C, iH, iW = inputs.size()

 

oH = iH // scale

 

oW = iW // scale

 

 

 

output = inputs.view(N, C, oH, scale, oW, scale)

 

output = output.permute(0, 1, 3, 5, 4, 2).contiguous()

 

return output.view(N, -1, oH, oW)

 

 

 

# Same as depthtospace of Tensorflow

 

def shuffle_up_tf(inputs, scale):

 

N, C, H, W = inputs.size()

 

inputs = inputs.view(N, scale ** 2, -1, H, W).transpose(2, 1).contiguous()

 

inputs = inputs.view(N, C, H, W)

 

return F.pixel_shuffle(inputs, scale)

 

 

 

# Inv DepthtoSpace (should be same as spacetodepth of Tensorflow, not test)

 

def shuffle_down_tf(inputs, scale):

 

output = shuffle_down_torch(inputs, scale)

 

N, C, H, W = output.size()

 

output = output.view(N, -1, scale * scale, H, W)

 

output = output.transpose(2, 1).contiguous().view(N, -1, H, W)

 

return output

 

参考文献
CVPR2016. Real Time Single Image and Video Super Resolution using an Efficient Sub-Pixel Convolutional Neural Network.
ECCV2018. Fast and Efficient Image Quality Enhancement via Desubpixel Convolutional Neural Network.
ICCV2019 Toward Real-World Single Image Super-Resolution A New Benchmark and A New Model.
Arxiv2019. Non-discriminative data or weak model? On the relative importance of data and model resolution.
Arxiv2020. TResNet: High Performance GPU-Dedicated Architecture.

Be First to Comment

发表评论

电子邮件地址不会被公开。 必填项已用*标注