HINet: Half Instance Normalization Network for Image Restoration
论文:HINet: Half Instance Normalization Network for Image Restoration
Abstract:
在本文中,我们探讨了实例归一化在低级视觉任务中的作用。 具体来说,我们提出了一个新颖的块:半实例归一化块(HIN 块),以提高图像恢复网络的性能。基于HINBlock,我们设计了一个简单而强大的多级网络,名为HINet,它由两个子网络组成。 在HIN Block的帮助下,HINet在各种图像恢复任务上超越了最先进的技术(SOTA)。对于图像去噪,我们在 SIDD 数据集上的 PSNR 超过了它 0.11dB 和 0.28dB,其乘法累加器运算 (MAC) 仅为 7.5% 和 30%,加速分别为 6.8 倍和 2.9 倍。对于图像去模糊,我们在 REDS 和 GoPro 数据集上获得了相当的性能,其 MAC 为 22.5%,加速率为 3.3 倍。 对于图像去雨,我们在多个数据集的平均结果上超过了 0.3 dB,加速比为 1.4 倍。 借助 HINet,我们在 NTIRE 2021 图像去模糊挑战赛 - Track2 中获得了第一名。 JPEG 伪影,PSNR 为 29.70。 代码可在 https://github.com/megviimodel/HINet 获取。
Introduction
归一化广泛应用于高级计算机视觉任务:分类中的Batch Normalization和IBN,DETR中的Layer Normalization和FCOS中检测的GroupNorm。此外,实例标准化用于样式/域传输任务。 然而,简单地将归一化应用于低级计算机视觉问题可能不是最理想的。 例如,Batch Normalization 无法提高网络在超分辨率方面的性能。
在本文中,我们仔细地将实例归一化作为构建块来集成,以提高图像恢复任务中的网络性能。 具体来说,我们提出了半实例归一化块(HIN 块)。 基于 HIN Blocks,我们进一步提出了一个称为 HINet 的多级网络,它由两个子网络组成。 通过在每个子网络的编码器中堆叠HIN Block,扩大了每个尺度的感受野,并且也提高了特征的鲁棒性。 除了每个阶段的架构之外,我们还在两个阶段之间采用跨阶段特征融合和监督注意模块来丰富多尺度特征并分别促进性能增益。
与最先进的架构 MPRNet相比,HINet 在各种图像恢复任务上都超越了它。 对于图像去噪,我们在 SIDD数据集上的 PSNR 超过了它 0.11 dB 和 0.28 dB,其乘法累加器运算 (MAC) 仅为 7.5% 和 30%,加速分别为 6.8 倍和 2.9 倍。 对于图像去模糊,我们在 REDS和 GoPro 数据集上获得了相当的性能,其 MAC 为 22.5%,加速率为 3.3 倍。 对于图像去雨,我们在 [56] 之后的多个数据集的平均结果上,PSNR 超过了它 0.3 dB,加速率为 1.4 倍。 各种图像恢复任务的可视化结果如图1所示。此外,我们将HIN应用于各种模型和各种数据集,结果证明了HIN的泛化能力。 例如,在 HIN 的帮助下,DMPHN在 GoPro 数据集上的 PSNR 增加了 0.42 dB。
我们的贡献可概括如下:
我们仔细地将实例归一化作为构建块进行集成,并提出了半实例归一化块。 据我们所知,它是第一个在图像恢复任务中直接采用归一化并具有最先进性能的模型。
基于 HIN Block,我们设计了一种用于图像恢复任务的多级架构 HINet,与 SOTA 方法相比,以更少的 MAC 和推理时间实现了最先进的性能。
进行了大量的实验来证明我们提出的 HIN 块和 HINet 的有效性。 在 HIN Block 和 HINet 的帮助下,我们在 NTIRE 2021 图像去模糊挑战赛 - Track2 上获得了第一名。 JPEG 伪影 ,PSNR 为 29.70。
Related Work
Normalization in low-level computer vision tasks:
标准化已成为高级计算机视觉任务的重要组成部分(尤其是批量标准化),但很少用于低级计算机视觉任务。 [31] 通过删除批量归一化修改了 ResBlock [15],因为他们在去模糊中使用大小为 2 的小批量训练模型。 [26]删除了超分辨率中的批量归一化,批量归一化消除了网络的范围灵活性。 由于[51]中讨论的图像恢复任务和批量归一化(BN)之间的不协调,图像恢复任务通常使用小图像块和小小批量大小来训练网络,这导致BN的统计不稳定。而且,图像恢复任务是每图像密集像素值预测任务,具有尺度敏感性。 而 BN 通常在规模不敏感的任务中很有帮助。
除了上述之外,还提出了实例归一化来取代[41]中的批量归一化,以提高风格迁移任务的性能。[16]证明实例归一化在某种程度上是低级特征的归一化。他们通过将样式图像的实例归一化中的通道统计数据与内容图像对齐,提出了对样式迁移任务的自适应实例归一化。基于[16],[20]采用自适应实例归一化作为正则化器来构建降噪器,并将从合成噪声数据中学到的知识转移到真实噪声降噪器。 与[20]不同,我们将实例归一化扩展为一种特征增强方法,并将其直接应用于图像恢复任务,而无需迁移学习。
Architectures for Image Restoration
单阶段方法广泛应用于图像恢复任务中,这些方法通常通过复杂的网络结构来提高网络容量。 多阶段方法将复杂的图像恢复任务分解为更小的更容易的子任务,在每个阶段都采用轻量级子网络。[11]将成熟的高斯-拉普拉斯图像金字塔分解技术引入到神经网络中,并使用相对较浅的网络来处理每个金字塔级别的学习问题。[35]通过重复展开浅层ResNet提出了渐进式循环网络,并引入循环层来利用跨阶段深层特征的依赖性。[58]提出了一种深堆叠分层多补丁网络。 每个级别侧重于不同尺度的模糊,较精细的级别将其残留图像贡献给较粗糙的级别。[55](Multi-stage progressive image restoration)提出了一种多级渐进图像恢复架构,其中有两个编码器-解码器子网络和一个原始分辨率子网络。 [55]还提出了每两个阶段之间的监督注意模块(SAM)和跨阶段特征融合(CSFF)模块,以丰富下一阶段的特征。 我们的模型还使用这两个模块来促进显着的性能增益,并使用两个简单的 U-Net作为子网络。
Approach
在本节中,我们将在以下小节中提供有关 HINet 和 HIN Block 的更详细说明。 具体来说,我们在3.1中引入HINet,在3.2中引入HIN Block。
HINet
我们提出的半实例归一化网络(HINet)的架构如图 2 所示。HINet 由两个子网络组成,每个子网络都是一个 U-Net。对于每个阶段的U-Net,我们使用一个3×3卷积层来提取初始特征。 然后,这些特征被输入到具有四次下采样和上采样的编码器-解码器架构中。我们使用内核大小等于 4 的卷积进行下采样,并使用内核大小等于 2 的转置卷积进行上采样。在编码器组件中,我们设计了半实例归一化块来提取每个尺度的特征,并在下采样时将特征通道加倍。在解码器组件中,我们使用ResBlocks来提取高级特征,并融合来自编码器组件的特征以补偿重采样造成的信息损失。对于 ResBlock,我们使用负斜率等于 0.2 的leaky ReLU,并删除批量归一化。 最后,我们使用一个 3 × 3 卷积得到重建图像的残差输出。
我们使用跨阶段特征融合(CSFF)模块和监督注意模块(SAM)来连接两个子网络,这两个模块来自[55]。 对于CSFF模块,我们使用3×3卷积将一个阶段的特征转换到下一阶段进行聚合,这有助于丰富下一阶段的多尺度特征。对于SAM,我们将原始模块中的1×1卷积替换为3×3卷积,并在每个卷积中添加偏差。 通过引入 SAM,当前阶段的有用特征可以传播到下一阶段,而信息量较少的特征将被注意力掩模抑制。
除了网络架构之外,我们还使用峰值信噪比(PSNR)作为损失函数的度量,即PSNR损失。设 X i ∈ R N × C × H × W X_{i}\in\mathbb{R}^{N\times C\times H\times W} Xi∈RN×C×H×W表示子网络 i 的输入,其中 N 是数据的批量大小,C 是通道数,H 和 W 是空间大小。 R i ∈ R N × C × H × W R_{i}\in\mathbb{R}^{N\times C\times H\times W} Ri∈RN×C×H×W表示子网络 i 的最终预测, Y i ∈ R N × C × H × W Y_{i}\in\mathbb{R}^{N\times C\times H\times W} Yi∈RN×C×H×W是每个阶段的真实情况。然后我们对HINet进行端到端优化如下:
L o s s = − ∑ i = 1 2 P S N R ( ( R i + X i ) , Y ) \begin{aligned}Loss=-\sum_{i=1}^2PSNR((R_i+X_i),Y)\end{aligned} Loss=−i=1∑2PSNR((Ri+Xi),Y)
Half Instance Normalization Block
由于小批量中小图像块的方差差异很大以及训练和测试的不同公式,BN在低级任务中并不常用。相反,实例归一化 (IN) 在训练和推理中保持相同的归一化过程一致。此外,IN在不受batch维度影响的情况下重新校准特征的均值和方差,比BN能够保留更多的尺度信息。 我们使用 IN 构建半实例归一化块(HIN 块)。通过引入HIN块,提高了HINet的建模能力(如图4所示)。 此外,IN引入的额外参数和计算成本可以忽略不计。
如图3a所示。 HIN 块首先采用输入特征 F i n ∈ R C i n × H × W F_{in}\in\mathbb{R}^{C_{in}\times H\times W} Fin∈RCin×H×W并通过 3 × 3 卷积生成中间特征 F m i d ∈ R C o u t × H × W F_{mid}\in\mathbb{R}^{C_{out}\times H\times W} Fmid∈RCout×H×W,其中 C i n / C o u t C_{in}/C_{out} Cin/Cout是 HIN 块的输入/输出通道数。 然后,特征 F m i d F_{mid} Fmid被分为两部分 ( F m i d 1 / F m i d 2 ∈ R C o u t / 2 × H × W ) (F_{mid_{1}}/F_{mid_{2}} \in \mathbb{R}^{C_{out}/2\times H\times W}) (Fmid1/Fmid2∈RCout/2×H×W)。第一部分 F m i d 1 F_{mid_{1}} Fmid1 通过可学习仿射参数的 IN 进行归一化,然后在通道维度中与 F m i d 2 F_{mid_{2}} Fmid2连接。HIN 块在一半通道上使用 IN,并通过另一半通道保留上下文信息。 后期的实验也会证明这种设计对于网络浅层的特征更加友好。经过concat操作后,将特征传递给一个3×3卷积层和两个leaky ReLU层,得到残差特征 R o u t ∈ R C o u t × H × W R_{out}\in\mathbb{R}^{C_{out}\times H\times W} Rout∈RCout×H×W,如图3a所示。 最后,HIN 通过添加具有快捷特征(1×1 卷积后获得)的 R o u t R_{out} Rout来块输出 F o u t F_{out} Fout。
Conclusion
在这项工作中,我们在图像恢复任务中重用归一化。 具体来说,我们将实例归一化引入到残差块中,并设计一个有效且高效的块:半实例归一化块(HIN块)。 在HIN Block中,我们对一半的中间特征应用Instance Normalization,同时保留内容信息。 基于HIN Block,我们进一步提出了一个称为HINet的多级网络。 在每个阶段之间,我们使用跨阶段的特征融合和注意力引导图[55]来缓解信息流动并增强多尺度特征表达。 我们提出的 HINet 在各种图像恢复任务上都超越了 SOTA。 此外,通过使用 HINet,我们在 NTIRE 2021 图像去模糊挑战赛 - Track2 中获得了第一名。 JPEG 伪像。