【全队项目】从GAN到ESRGAN的超分辨率处理

直入主题

ESRGAN的全名叫 Enhanced Super-Resolution Generative Adversarial Networks,发表于ECCV2018,它是基于SRGAN改进而来到,相比于SRGAN它在三个方面进行了改进

改进了网络结构,对抗损失,感知损失
引入**Residual-in-Residu Dense Block**(RRDB)
使用激活前的VGG特征来改善感知损失

在开始讲这个ESRGAN的具体实现之前,我先来看一下他和他的前辈SRGAN的对比效果:

我们可以从上图看出,

ESRGAN在锐度和边缘信息上优于SRGAN,且去除了"伪影"
从PI和PMSE两个指标来看,ESRGAN也可以当之无愧地称得上是超分辨率复原任务中的the State-of-the-Art

庖丁解牛

RRDB,对residual blocks的改进:

我们可以看出这个残差块是很传统的Conv-BN-relu-Conv-BN的结构,而作者在文章中是这么说到的:

We empirically observe that BN layers tend to bring artifacts. These artifacts, namely BN artifacts, occasionally appear among iterations and different settings, violating the needs for a stable performance over training. In this section, we present that the network depth, BN position, training dataset and training loss have impact on the occurrence of BN artifacts and show corresponding visual

什么意思呢?就是说作者认为SRGAN之所以会产生伪影,就是因为使用了Batchnormalization,所以作者做出了去除BN的改进

而且我们再来看,SRGAN的残差块是顺序连接的,而作者可能哎,受denseNet的启发,他就把这些残差块用密集连接的方式连在一起.那么他的生成器里的特征提取部分最终变成了这样子:

既然我们知道他的网络结构是这样子设计的,那么他的实现其实就很简单:

复制代码

class RRDB_Net(nn.Module):
    def __init__(self, in_nc, out_nc, nf, nb, gc=32, upscale=4, norm_type=None, act_type='leakyrelu', \
            mode='CNA', res_scale=1, upsample_mode='upconv'):
        super(RRDB_Net, self).__init__()
        n_upscale = int(math.log(upscale, 2))
        if upscale == 3:
            n_upscale = 1

        fea_conv = B.conv_block(in_nc, nf, kernel_size=3, norm_type=None, act_type=None)
        rb_blocks = [B.RRDB(nf, kernel_size=3, gc=32, stride=1, bias=True, pad_type='zero', \
            norm_type=norm_type, act_type=act_type, mode='CNA') for _ in range(nb)]
        LR_conv = B.conv_block(nf, nf, kernel_size=3, norm_type=norm_type, act_type=None, mode=mode)

        if upsample_mode == 'upconv':
            upsample_block = B.upconv_blcok
        elif upsample_mode == 'pixelshuffle':
            upsample_block = B.pixelshuffle_block
        else:
            raise NotImplementedError('upsample mode [%s] is not found' % upsample_mode)
        if upscale == 3:
            upsampler = upsample_block(nf, nf, 3, act_type=act_type)
        else:
            upsampler = [upsample_block(nf, nf, act_type=act_type) for _ in range(n_upscale)]
        HR_conv0 = B.conv_block(nf, nf, kernel_size=3, norm_type=None, act_type=act_type)
        HR_conv1 = B.conv_block(nf, out_nc, kernel_size=3, norm_type=None, act_type=None)

        self.model = B.sequential(fea_conv, B.ShortcutBlock(B.sequential(*rb_blocks, LR_conv)),\
            *upsampler, HR_conv0, HR_conv1)

    def forward(self, x):
        x = self.model(x)

对损失函数的改进

说到损失函数啊,我们之前在SRGAN的文章里我也介绍过,这个判别器它判断的是你输入的图片是"真的"高清图像,还是"假的"高清图像,而且作者他就提出一种新的思考模式,就是说我的判别器是来估计真实图像相对来说比fake图像更逼真的概率。

怎么来理解这句话呢?

具体而言，作者把标准的判别器换成Relativistic average Discriminator（RaD），所以判别器的损失函数定义为：

对应的生成器的对抗损失函数为：

求MSE的操作是通过对 $mini-batch$ 中的所有数据求平均得到的，$$x_f$$是原始低分辨图像经过生成器以后的图像,由于对抗的损失包含了$$x_r$$和$$x_f$$,所以生成器受益于对抗训练中的生成数据和实际数据的梯度，这种调整会使得网络学习到更尖锐的边缘和更细节的纹理。

因为作者放放出来的代码比较emmmmmm,没找到具体的训练部分哎......

对感知损失的改进

我们之前看SRGAN的时候看到哎,它是用来一个训练好的VGG16来给出超分辨率复原所需要的特征,作者通过对损失域的研究发现,激活前的特征，这样会克服两个缺点。

Perceptual loss is previously defined on the activation layers of a pre-trained deep network, where the distance between two activated features is minimized. Contrary to the convention, we propose to use features before the activation layers, which will overcome two drawbacks of the original design.

激活后的特征是非常稀疏的，特别是在很深的网络中。这种稀疏的激活提供的监督效果是很弱的，会造成性能低下；
使用激活后的特征会导致重建图像与GT的亮度不一致。

与此同时,作者还在loss函数中加入了$$L_1=E_{xi}||G(x_i)-y||_1$$,也就是$$L_1$$损失,最终损失函数由三部分组成:

网络插值

为了平衡感知质量和PSNR等评价值，作者提出了一个灵活且有效的方法---网络插值。具体而言，作者首先基于PSNR方法训练的得到的网络G_PSNR，然后再用基于GAN的网络G_GAN进行整合

它的具体实现就是:

复制代码

for k, v_PSNR in net_PSNR.items():
    v_ESRGAN = net_ESRGAN[k]
    net_interp[k] = (1 - alpha) * v_PSNR + alpha * v_ESRGAN

具体实现部分

ESRGAN可以实现放大4倍的效果

首先要训练一个基于PSRN指标的模型,如何根据这个模型的权重进行生成器的初始化.

作者使用了Adam作为优化器(β1 = 0.9, β2 = 0.999.)进行交替训练.

生成器有16个residual block和23个RRDB