论文链接:https://arxiv.org/abs/2112.04491
图像恢复任务,是指将受损(如带噪声/模糊)的图像恢复为清晰图像。这在日常生活中广泛出现,如手机拍摄的照片通常需要经过图像恢复算法对其进行去噪/去模糊等一系列处理之后,再显示给用户。
随着人工智能的崛起,深度学习也占领了图像恢复领域的高地:近年的模型如HINet[1], MPRNet[2], Restormer[3] 等等均在该领域取得了领先的性能。
本文首先介绍基于深度学习的图像恢复任务的流程背景,然后重新审视模型中常见的全局信息聚合操作并揭示现有方法中被大家广泛忽视的问题:训练和测试的不一致性。最后,我们会介绍一个简单的解决方法——测试时局部转换器(Test-time Local Converter, 简称TLC),通过缓解不一致性问题来提升模型性能。
背景:图像恢复任务流程
图1. 图像恢复任务训练/测试过程示意图
在实际应用中,用于图像恢复任务(如去模糊)的数据集通常由高分辨率图像组成。由于数据增广的需要和GPU显存的限制,通常的做法是用从高分辨率图像中裁剪的小块来训练模型。以GoPro数据集为例,MPRNet[2]在训练时使用从原图中裁剪出的256x256的小块(如图1. 上半部分所示),而在推理(inference)过程中,为了利用全图信息,训练好的模型会直接复原720x1280的高分辨率图像(如图1. 下半部分所示)。
所以,模型的输入在训练/推理的过程中存在不一致:即训练时是图像的局部区域,而推理时则是整个图像。后续我们将展示这种不一致性对利用全局信息的模型产生的负面影响。
重新审视全局信息聚合
图2. 全局信息聚合的例子
全局信息聚合操作指的是沿着输入特征的整个空间维度聚合信息。以全局平均池化(global average pooling)操作为代表的全局操作被广泛用于目前性能领先的图像复原模型中。例如,HINet[1]中的Instance Normalization使用全局的均值和方差对特征沿着整个空间维度执行全局归一化,如图2. 左半部分。MPRNet[2]、FFA-Net[4]等使用了通道注意力机制(如SE[5]模块)。该模块使用全局平均池化操作聚合全局信息用于选择性地强调通道特征,如图2. 右半部分。
图3. 全局操作在训练/推理阶段的不一致性
然而, 全局操作在训练和测试阶段的表现是不一致的:它们在训练中是基于裁剪图像的部分信息,而在测试时是基于全分辨率图像的所有信息,这可能导致特征的分布完全不一致。
我们进一步分析了全局平均池化操作的结果(即特征图的均值),发现它的分布在训练/测试过程中差异较大。具体而言,我们将MPRNet[2]中SE[5]模块的全局平均池化结果分布进行了可视化,如图4. 所示。其中,绿色表示基于裁切图像的数据分布,而蓝色表示基于全分辨率图像的数据分布。可以看到,绿色到蓝色所展示的信息分布差异较大。这种训练和推理中全局信息分布的变化会对模型的性能产生负面影响。
图4. MPRNet特征全局池化后分布可视化
我们提出了一个简单的测试时方法 来减少训练和测试的不一致性。使用了本文方法后的数据分布如图4. 右侧红色所示,接下来将具体介绍本文的方法。
测试时局部转换器
(Test-time Local Converter, 简称 TLC)
我们提出了一个名为测试时局部转换器Test-time Local Converter(TLC)的推理方案,以减小训练和推理之间的全局信息的分布差异。如图5.b. 所示,在推理阶段,我们的TLC直接改变了特征层面的信息聚合的区域范围,将全局操作(例如全局平均池化)仅在推理过程中被转换为局部操作。
这样的好处是,测试阶段每个点的统计量分布与训练时的分布相近(均基于局部区域)。具体来说,全局操作的输入特征被切成相同大小的重叠窗口,然后对每个重叠窗口独立地进行信息聚合操作。从理论时间复杂度的角度,TLC可以借助前缀和技巧从而避免额外的计算复杂度;从推理时间角度来说,经过优化的TLC仅会带来1%~5%左右的时间增长,却能带来显著的性能提升。
该技术是通用的,它可以应用在多种全局操作之上,并且无需任何额外训练或微调,在仅增加少量计算开销的情况下显著提高各种模块(如SE、IN)的性能。
与切块测试的比较
图6. 与切块测试的比较
还有一种简单的方法是将图像分成多个小块进行独立推理可以缓解不一致性问题。然而,由于边界像素不能利用区域之外的邻近像素信息,该方法可能会在每个小块周围引入边界伪影。
如图6. 右上角所示,直接对图像进行分快预测引入了一条明显的垂直分割线,严重损害了图像质量。如图6. 右下角所示,我们的方法生成的图像更加清晰且无伪影。
实验结果
我们展示部分实验结果。不需要重新训练和精调,将TLC直接应用于HINet[1], MPRNet[2], Restormer[3](去模糊), FFANet[4](去雾)等模型,称为HINet-Local, MPRNet-Local, Restormer-Local 以及 FFANet-Local,其客观指标均有明显增长。
主观指标如图7. 图8. 所示,应用TLC之后,图片主观质量有明显提升。
图8. 去雾可视化效果
此外,TLC也协助我们在图像去模糊/去噪/双目超分辨率等领域取得最先进的性能[6],[7]。
更多实验结果及细节请见我们的论文,代码也已开源,欢迎试用。
来源:旷视研究院
作者:R
专栏文章推荐
- ECCV 2022 | 图像恢复的简单基线
- ECCV 2022 | 视频插帧中的实时中间流估计
- 沈向洋对话张祥雨:重新审视CNN中大卷积核问题
- ECCV 2022 Oral | 基于EM算法从视频中生成真实光流数据集
欢迎关注旷视研究院极术社区专栏,定期更新最新旷视研究院成果。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。
加入旷视:career@megvii.com