消除图像复原中的“misalignment”，性能大幅提升

arXiv:https://arxiv.org/pdf/2112.04491.pdf
code:https://github.com/megvii-research/tlsc

"misalignment"之处必有大发现。CNN训练与测试图像分辨率不一致引申出了FixResNet与FixEfficientNet；目标检测中Head在训练与测试过程中的不一致同样引申了不少优秀工作。那么图像复原领域有没有这类现象呢？那必须有哇，所以看看这篇paper咯。

Abstract

全局空域统计(即沿完整空间维度进行信息聚合)已被广泛用于图像复原方案中。比如，HINet所使用InstanceNorm的均值与方差，MPRNet所使用的SE的全局均值池化。本文首次表明：训练/测试阶段的基于图像块/完整图像特征的统计聚合计算差异会导致不同的分布，进而导致图像复原的性能下降(该现象被广泛忽视了)。

为解决该问题，我们提出一种简单的TLSC(Test-time Local Statistics Converter)方案，它在测试阶段的区域统计聚合操作由全局替换为局部。无需重新练或微调，所提方案可以大幅提升已有图像复原方案的性能。

实验表明：采用TLSC对SE进行扩展，MPRNet在GoPro数据集上的性能提升了0.65dB(参见下图)，达到了33.31dB，以0.6dB超越此前最佳结果。此外，我们将TLSC用于High-level任务(如语义分割)后同样取得了有竞争力的结果。定量与定性实验结果均表明：TLSC能够以轻微的额外计算消耗取得大幅性能增益。

StartPoint

在正式引出本文方案之前，我们先对图像复原流水线及其产生的统计不一致问题进行描述；然后再引出本文的解决方案TLSC。

上图给出了图像复原方案在训练与测试阶段的处理流水线示意图。可以看到：

在训练阶段，输入图像先crop成图像块，然后再送入网络中进行全局统计信息聚合；
在测试阶段，则由不同的处理方式，如上图b的全图输入以及上图c的分块输入。

在训练阶段，参数以块统计方式进行优化；而在测试阶段则是以全图统计方式推理。训练与测试阶段的不同全局统计聚合计算方式就导致了"misalignment"，即出现了统计分布不一致现象。图像复原方案难以针对统计分布的严重改变进行自适应调整，进而导致性能退化。将图像拆分为图像块进行推理可以避免上述不一致现象，但会导致新的边界伪影问题。

Test-time Local Statistics Converter

基于上述不一致现象以及引发的问题，我们总结出一个实用方案应该考虑如下几个条件：

为确保广泛应用，它需要在无需重训练或微调仍然有效，即应当适用于已训练模型；
训练与测试阶段的统计分布应当匹配，否则性能退化就会发生；
为保证图像质量，不能引入明显的边界伪影；
不能导致计算瓶颈问题，否则实际场景会忽略该方案。

本文提出一种新的方案：(1)它以全图作为输入以避免边界伪影；(2) 特征图中的每个像素仅在局部窗口计算各自统计信息，故可确保训练/测试时的统计一致性。

为简单起见，上式并未考虑边界像素情形。实际上，我们可以采用两步方式进行实现：

对于非边界像素，以滑动窗口形式对每个像素进行局部统计信息计算
对于边界像素，通过复制的方式填充。

总而言之，该方案可以总结为：将统计聚合操作由全局转换为局部。故将其命名为TLSC。

Extending TLSC

接下来，我们将介绍如何将TLSC嵌入到SE与IN模块中，其他GN、CBAM类似可得。

Experiments

我们先来看关于TLSC使用前后的定量与定性分析，并由此得出超参的选择方式。

上图给出了训练与测试时的统计分布不一致现象的示意图，而TLSC则可以解决该问题，使得训练与测试时的统计分布相一致。为进一步证实该不一致会导致性能退化，我们以UNet、UNet-IN、UNet-SE进行了实验对比，见下图。

上图对比了不同块尺寸时不同模型的性能，从中可以看到：UNet-IN与UNet-SE的性能会随块尺寸增大出现性能下降。我们猜测：测试时更大的块可以利用更多的信息，故UNet取得了更佳性能。然而，UNet-SE与UNet-IN并未从更大测试块中受益。这是由训练与测试时的统计不一致导致。

将图像拆分为块进行推理可以消除统计不一致现象，但会引入边界伪影问题(见下图)，进而影响图像质量。而使用全图进行测试会导致严重的性能下降，见上表。

上图与下表给出了使用TLSC前后的性能对比，可以看到：TLSC可以分别提升UNet-IN与UNet-SE的性能0.16/0.39dB，进而达到31.11/31.21dB，大幅超越了Baseline的30.86dB。这说明：TLSC可以通过减少统计分布漂移改善模型性能。

此外，上表还提供了TLSC在CBAM、GN以及GE方面的扩展性能对比，均可以看到显著的性能提升，这进一步验证了TLSC的可扩展能力。

局部窗口尺寸是TLSC的超参数，它用于控制局部统计聚合操作的范围。上表对比了不同训练尺寸与测试尺寸时的性能对比关系，可以看到：测试阶段的最优窗口尺寸并不完全相同，但比训练尺寸更大。我们猜测这是如下原因导致：大尺寸时的更多信息受益与统计不一致影响之间的均衡。此外，由于该方案不需要重训练，故它可以灵活的调整局部窗口尺寸。