• 问答
  • 技术
  • 实践
  • 资源
性能大幅提升!消除图像复原中的 “misalignment”
技术讨论

作者丨happy
来源丨AIWalker
编辑丨极市平台

论文链接:https://arxiv.org/pdf/2112.04491.pdf

代码链接:https://github.com/megvii-research/tlsc

"misalignment"之处必有大发现。CNN训练与测试图像分辨率不一致引申出了FixResNet与FixEfficientNet;目标检测中Head在训练与测试过程中的不一致同样引申了不少优秀工作。

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的分块输入。

特征$X \in R^{H \times W}$全局统计聚合可以描述如下:

$$
\Phi(X, f) = \frac{1}{HW} \sum{p=1}^H \sum{q=1}^W f(X_{p,q})
$$

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

Test-time Local Statistics Converter

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

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

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

具体来说将,特征$X \in R^{H \times W}$的局部统计聚合描述如下:

$$
\Psi(X, f)_{i,j} = \frac{1}{K_h K_w} \sum_p \sumq f(X{X_{p,q}})
$$

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

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

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

Extending TLSC

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

Extending TLSC to SE Block. SE首先对输入特征进行全局空域信息统计并挤压为通道,可表示为$\Phi(X^{c}, id), \forall c \in [C], id(t)=t, \forall t \in R$;然后采用一个MLP计算通道注意力并对原始特征重加权。

由于统计分布漂移,全局空域信息挤压是一种次优方案。为解决该问题,我们采用TLSC对SE进行扩展:将$\Phi(X^{c}, id) $替换为$\Psi(X^{c}, id)$,MLP部分则替换为$1\times 1$卷积即可。此时,SE的通道注意力就变成了元素注意力。

Extending TLSC to IN Block. 对于特征$X \in R^{H \times W}$而言,规范化特征Y计算如下

$$
Y = \frac{X - \mu}{\sigma} \
\mu = \Phi(X, id) \
\sigma^2 = \Phi(X, sq) - \mu^2
$$

为简单起见,我们忽视掉了IN中的$\gamma, \beta$参数。在推理阶段,我们将$\Phi(X, id), \Phi(X, sq)$替换为如下局部版本:

$$
\mu{local} = \Psi(X, id) \
\sigma
{local}^2 = \Psi(X, sq) - \mu_{local}^2
$$

此时,规范化特征Y计算方式如下:

$$
Y{local} = \frac{X - \mu{local}}{\sigma_{local}}
$$

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

Applications

为进一步验证所提方案的有效性,我们在图像复原与语义分割任务上进行了验证。

上表给出了图像去模糊、图像去雨、图像去雾三种不同复原任务上的性能对比。可以看到:

  • 对于图像去模糊,HINet、MPRNet与TLSC相结合时均可看到显著的性能提升,并取得了新的SOTA结果;
  • 对于图像去雨,SPDNet与TLSC相结合时取得了0.2dB指标提升;
  • 对于图像去雾,FFANet与TSLC相结合时取得了0.42dB(户外数据)指标提升;
  • 上述不同任务的性能提升表明:TLSC可以有效提升不同复原任务中带全局统计聚合操作模型的性能

上表给出了语义分割任务中采用TLSC前后的性能对比,可以看到:TLSC可以提升所有模型性能超1.11%!,代价为少于2%的额外计算量。这说明:TLSC对于分割任务同样非常有效**。

  • 0
  • 0
  • 252
收藏
暂无评论