Progressive Feature Fusion Networkfor Realistic Image Dehazing笔记及代码

Progressive Feature Fusion Networkfor Realistic Image Dehazing笔记及代码

来源

Asian Conference on Computer Vision2018

CCF c类会议

摘要

• 提出一个端到端的网络结构
• 提出了一种基于渐进特征融合的U-Net类编码器-解码器网络，直接学习从观测到的模糊图像到无模糊图像的高度非线性变换函数（不依赖大气散射模型）。

提出的方法

本文提出了基于U-Net的渐进式网络，它由编码器、特征转换器、解码器三个模块组成，如下图所示。

编码器由5个卷积层组成，每一个卷积层后跟着一层ReLU激活函数层。D0用于从原始观测到的模糊图像中聚合一个相对较大的局部接受域上的信息特征。然后，以下四层依次执行向下采样卷积操作，以金字塔比例对图像的信息进行编码。如果输入的图像大小为w * h * c，编码器D4的特征图大小为（w/16)* (h/16) * c。

特征变换模块，由基于残差的子网络组成。为了平衡计算效率和GPU内存的使用，在这个模块中，我们使用了18个残差块进行特征学习，卷积层的通道均为256，它们与来自编码器模块的特征图D4的通道相同。步长始终保持在1。

解码器模块由四个反卷积层和一个卷积层组成。解码器的反卷积层按顺序恢复图像的结构细节，每一层后面跟着ReLU层。为了最大限度地实现多层信息流，保证更好的收敛性，在编码器和解码器对应的不同层次的层之间采用了跳跃连接。特征变换模块的输入和输出之间采用全局快捷连接，如图2所示。

因此，转置特征图U0的维数为×16，与×10相同。对U0进一步应用卷积运算，生成最后清晰图像J(x)。在这里，对于这个卷积层，核大小k为3；步长为1；且通道与J相同。

网络逐步对编码器和解码器之间的空间金字塔映射进行特征融合，从而最大限度地从反卷积层的输入中保留结构细节，进一步使去模糊网络更具输入自适应能力。

实验

数据集

• NTIRE2018 Image Dehazing Dataset. 分为I-HAZE和O-HAZE。
• RESIDE

在本文中，同时使用I-HAZE和O-HAZE训练图像进行网络训练，总共有80个场景。基于原数据集进行数据增强，扩充数据集。将数据集剪切成512 * 512尺寸的数据集。Adam作为优化器，学习率设为0.0001。以恢复的清晰图像和无雾地面真实之间的均方误差(MSE)作为我们的客观损失。batch-size设为32。

在本文修复PFFNet的体系结构之前，做了几个消融实验。在特征转换模块中实验了四种不同的块大小：{6、12、18、24}。测试性能如图4-左所示。增加残余块的大小将提高测试性能。通过考虑性能与可用计算资源之间的平衡，本文最终采用了具有***18个残差块***的特征变换模块。

代码

残差块

每两层进行一次跳跃连接，便构成了残差块。残差块源自何凯明的ResNet，残差的思想有效的缓解了神经网络的梯度消失现象。我们通过跨层连接后面的网络层，跳过这2个卷积运算，将输入直接加在最后的ReLU激活函数前。这样要求2个卷积层的输出与输入形状一样，从而可以相加。简单的说就是 out = layer(input) + input

代码：

class ResidualBlock(torch.nn.Module):
    def __init__(self, channels):
        super(ResidualBlock, self).__init__()
        self.conv1 = ConvLayer(channels, channels, kernel_size=3, stride=1)
        self.conv2 = ConvLayer(channels, channels, kernel_size=3, stride=1)
        self.relu = nn.PReLU()

    def forward(self, x):
        residual = x
        out = self.relu(self.conv1(x))
        out = self.conv2(out) * 0.1
        out = torch.add(out, residual)
        return out

U-Net结构

编码器

输入的第一个卷积层是一个卷积核为11的卷积层，

x = self.relu(self.conv_input(x))

后面跟着四个卷积模块

res2x = self.relu(self.conv2x(x))
res4x = self.relu(self.conv4x(res2x))

res8x = self.relu(self.conv8x(res4x))
res16x = self.relu(self.conv16x(res8x))

特征转换器

特征转换器包含着数个残差块，res_blocks表示残差块的数量。

self.dehaze = nn.Sequential()
        for i in range(1, res_blocks):
            self.dehaze.add_module('res%d' % i, ResidualBlock(256))

操作如下：

res_dehaze = res16x
res16x = self.dehaze(res16x)
res16x = torch.add(res_dehaze, res16x)

解码器

解码器包含着四个反卷积层，有编码器四个卷积层相对应，对应进行特征的跳跃连接。

self.convd16x = UpsampleConvLayer(256, 128, kernel_size=3, stride=2)
self.convd8x = UpsampleConvLayer(128, 64, kernel_size=3, stride=2)
self.convd4x = UpsampleConvLayer(64, 32, kernel_size=3, stride=2)
self.convd2x = UpsampleConvLayer(32, 16, kernel_size=3, stride=2)

操作如下:

res16x = self.relu(self.convd16x(res16x))
res16x = F.upsample(res16x, res8x.size()[2:], mode='bilinear')
res8x = torch.add(res16x, res8x)

res8x = self.relu(self.convd8x(res8x))
res8x = F.upsample(res8x, res4x.size()[2:], mode='bilinear')
res4x = torch.add(res8x, res4x)

res4x = self.relu(self.convd4x(res4x))
res4x = F.upsample(res4x, res2x.size()[2:], mode='bilinear')
res2x = torch.add(res4x, res2x)

res2x = self.relu(self.convd2x(res2x))
res2x = F.upsample(res2x, x.size()[2:], mode='bilinear')
x = torch.add(res2x, x)

最后跟一个步长为1的卷积层。

self.conv_output = ConvLayer(16, 3, kernel_size=3, stride=1)