基于窗剪枝的去局部模糊Transformer，提升像质同时加速推理

2023-07-11 05:07

在日常摄影中，由于曝光时间段内运动物体和静止背景之间的混叠，照片中会出现运动模糊。随着云台等稳像设备的发展，由于相机抖动所导致的全局运动模糊问题大大减少，图像中的局部运动模糊成为了新的挑战。现有的图像去模糊方法主要解决全局运动模糊问题，对图像中所有的区域都进行去模糊操作。这种全局去模糊的做法不仅会降低非模糊区域的清晰度，也浪费了许多计算资源，尤其是重建含有局部运动模糊的大分辨率的图像时，计算资源消耗巨大。

为了灵活、有效地、快速地恢复含有局部运动模糊的大分辨率图像，新加坡南洋理工大学 S-Lab 联合浙江大学提出了一个自适应的图像去局部运动模糊网络 LMD-ViT，和窗剪枝变换器 AdaWPT，将去模糊的重点放在局部模糊区域，对非模糊区域以窗口为单位剪枝，大大减少了计算量（FLOPs-66%），实现了推理速度两倍以上的提升，在去模糊性能上也并没有打折，甚至相较于全局去模糊 Transformer 提升 PSNR 0.24dB，实现了又快又好地去局部运动模糊。代码即将开源。

论文标题：

Adaptive Window Pruning for Efficient Local Motion Deblurring

作者单位：

NTU S-Lab，ZJU

论文链接：

https://arxiv.org/pdf/2306.14268.pdf

论文主页：

https://leiali.github.io/LMD-ViT_webpage/index.html

效果展示

下图展示了该文章所提出的 LMD-ViT 方法的关键过程的可视化结果。LMD-ViT 对不同分辨率的特征图进行不同粒度的窗剪枝操作，下图中白色模板覆盖的区域为保留的窗口，其余区域对应被剪枝的窗口，Block #4、6、8 对应不同网络模块中不同粒度的可视化窗剪枝结果。在 ReLoBlur 测试数据集上，全局去模糊方法（如：Uformer [1]）会使原本清晰的区域变形，而 LMD-ViT 在有效去局部模糊的同时，保护了清晰区域不受干扰。

下图展示了 LMD-ViT 的去局部运动模糊效果。

方法概览

2.1 自适应去除局部运动模糊网络 LMD-ViT

LMD-ViT 是一个 U 形网络，由若干编码阶段、瓶颈阶段和解码器阶段组成，编码阶段和对应的解码阶段相连。在网络的开始和结束处分别放置了一个投影层和输出层，将 RGB 图像提取为特征图或将特征图转换为 RGB 图像。网络的每个阶段都包含一系列不同分辨率的自适应窗剪枝变换器（AdaWPT）和下采样/上采样层，不同分辨率的 AdaWPT 在特征图上进行不同粒度的剪枝。

2.2 自适应窗剪枝变换器（AdaWPT）

作为一个关键组成部分， AdaWPT 包括第一 AdaWPT 模块（AdaWPT-F）和若干后 AdaWPT 模块（AdaWPT-P）。在训练过程中，每个 AdaWPT-F 包括一个置信度预测器（Confidence Predictor）、一个决策层（Decision Layer）、一个特征转换/反转模块和若干 Transformer 层，Transformer 层包括基于窗口的多头注意力层（W-MSA）、基于窗口的局部增强前馈层（W-LeFF）和归一化层（LN）。

推理过程较训练过程多引入了一个窗口合并操作（Window Compound）。AdaWPT-F 通过 Confidence Predictor 预测各窗口的模糊置信度，经过 Decision Layer 决定哪些窗需要保留，哪些需要丢弃。为了节省计算资源，AdaWPT-P 按照 AdaWPT-F 提供的决策进行剪枝，不需要再次预测模糊置信度或做决定。此外，在剪枝操作各前后插入了一个特征转换/反转块，以促进图像特征的相互作用。

在推理过程中，Decision Layer 使用 Gumble-Softmax 机制，模糊置信度高的窗口有更大可能性被置为 1，反之则被置为 0；被置为 0 的窗口将被剪枝，只有未剪枝的窗口进入 Transformer 层，依次经历 W-MSA、W-LeFF 和 LN 等操作，生成去模糊的特征窗口；在 Transformer 层后，被剪枝的窗口和去模糊的特征窗口重组成符合空间顺序的特征图。这种基于窗剪枝的推理策略在去除局部运动模糊的同时，保证了清晰区域不失真，且大大降低了计算成本。

在训练过程中，Decision Layer 使用 Softmax 机制，并配合自定义阈值使用；为了保证并行训练和反向传播，被置为 0 的窗口和没有被剪枝的窗口一起进入 Transformer 层。

2.3 局部运动模糊区域标定

为了更好地训练 Confidence Predictor，该文章手工标定了 ReLoBlur 训练数据集中的模糊区域，和 LBFMF [2] 用算法标注的模糊区域相比，我们手工标注的模糊区域不含有空洞和噪声，且更符合人眼对运动模糊区域的判断，实验也证明了该文章提供的标注模版提升了去模糊网络的训练效果。

2.4 更多结果展示

LMD-ViT 对不同分辨率的特征图进行不同粒度的窗剪枝操作，下图中白色模板覆盖的区域为保留的窗口，其余区域对应被剪枝的窗口。

LMD-ViT 较 CNN 类的去局部运动模糊方法，实现了更好的去模糊效果；较 Transformer 类方法，不仅提高了去模糊效果，还具有更快的推理速度和更少的 FLOPs。

LMD-ViT 在对局部进行去模糊的同时，不对清晰区域产生形变、模糊等降质影响。

LMD-ViT 也可以解决全局模糊问题，当输入整张图都模糊的图片，网络不对窗口剪枝，所有的窗口都进行去模糊操作。

在实拍图像去模糊效果的用户调研中，LMD-ViT 收到了最多数用户的认可。

总体而言，该文章提出的基于 Transformer 的自适应窗剪枝网络，和一系列针对大分辨率图像的去除局部运动模糊策略，有效、快速地去除了图像中的局部运动模糊，使图像获得了明显的主客观像质提升；该文章对模糊区域精细的标注，也大大提升了当下去局部运动模糊方法的训练效果。整体上，该文章灵活地解决了图像去运动模糊问题，为工业界和学术界提供了较为实用的局部像质提升思路。