【ECCV2022】Improving Image Restoration by Revisiting Global Information Aggregation
代码:https://github.com/megvii-research/tlc
旷视研究院的工作,论文考虑了一个非常有趣的问题:由于GPU显存的限制,通常图像修复是从高分辨率图像中裁剪小块来训练模型。下图为例,当前方法在训练时使用从原图中裁剪出的256x256的小块,而在推理过程中,训练好的模型会直接复原720x1280的高分辨率图像。因此,训练/推理的过程中存在不一致(训练时是图像的局部区域,而推理时则是整个图像),对利用全局信息的模型产生的负面影响。
当前的一些方法,如 HINet 中的 Instance Normalization 使用全局的均值和方差对特征沿着整个空间维度执行全局归一化,如下图左侧;MPRNet、FFA-Net等使用了通道注意力机制(SE模块),使用全局池化操作聚合全局信息,如下图右侧。但是, 全局操作在训练和测试阶段的表现是不一致的:它们在训练中是基于裁剪图像的部分信息,而在测试时是基于全分辨率图像的所有信息,这可能导致特征的分布完全不一致。
为了解决上面的问题,作者提出了 Test-time Local Converter (TLC) ,以减小训练和推理之间的全局信息的分布差异。如下图所示,在推理阶段,TLC直接改变了特征层面的信息聚合的区域范围,将全局操作(例如全局平均池化)仅在推理过程中被转换为局部操作。
这样的好处是,测试阶段每个点的统计量分布与训练时的分布相近(均基于局部区域)。具体来说,全局操作的输入特征被切成相同大小的重叠窗口,然后对每个重叠窗口独立地进行信息聚合操作。经过优化的TLC仅会带来1%~5%左右的时间增长,却能带来显著的性能提升。
作者指出,窗口需要是重叠的,不然会引入边界伪影,如下图所示。
实验部分,不需要重新训练和精调,将TLC直接应用于HINet, MPRNet, Restormer(去模糊), FFANet(去雾)等模型,客观指标均有明显增长。