Dynamic Pre-training：实现高效、可扩展的一体化(All-in-one)图像恢复

题目： Dynamic Pre-training: Towards Efficient and Scalable All-in-One Image Restoration 作者： Akshay Dudhane, Omkar Thawakar, Syed Waqas Zamir等 论文链接：https://arxiv.org/abs/2404.02154 代码链接：https://github.com/akshaydudhane16/DyNet 来源：arxiv 2024 内容整理：周楚骎 一体化图像修复使用统一的模型处理不同类型的退化，而不是为每种退化使用特定于任务的非通用模型。使用同一模型处理多种退化问题的需求可能导致具有固定配置的高复杂性设计，这些设计缺乏对更高效替代方案的适应性。本文提出了 DyNet，一个动态的网络系列，以编码器-解码器风格设计，用于一体化图像恢复任务。本文的 DyNet 可以在其笨重和轻量级变体之间无缝切换，并且通过一轮训练为高效模型部署提供灵活性。这种无缝切换是通过本文的权重共享机制实现的，它构成了本文架构的核心，并促进了初始化模块权重的重用。此外，为了建立稳健的权重初始化，本文引入了一种动态预训练策略，该策略同时训练所提出的 DyNet 的变体，从而将 GPU 训练时间数减少 50%。为了解决预训练中所需的大规模数据集的不可用问题，本文策划了一个名为 Million-IRD 的高质量、高分辨率图像数据集，其中包含 2M 个图像样本。本文验证了 DyNet 在一体化环境中的图像去噪、去雨和去雾，与基线模型相比，GFlops 降低了 31.34%，参数减少了 56.75%，取得了最先进的结果。

目录

• 简介
• 方法
  • 动态网络 （DyNet） 架构
  • 动态预训练策略
  • Million-IRD：图像修复数据集
  • 数据后处理
• 实验
  • 实验细节
  • All-in-one设置下多重退化的比较
  • 单一退化的比较
  • 消融实验
• 总结

简介

图像恢复 （IR） 任务旨在改进低质量的输入图像。尽管 IR 取得了一些进展，但图像中存在的各种降解类型和严重程度仍然构成重大挑战。大多数现有方法隐式学习图像先验，需要针对不同的退化类型、水平和数据集进行单独的网络训练。此外，这些方法需要先验的图像退化知识，以便在测试期间进行有效的模型选择，因此缺乏通用性来满足各种退化的需求。

一体化恢复旨在使用单个统一模型恢复具有未知退化的图像。最近的进展，如AirNet和PromptIR，分别通过采用对比学习和隐式视觉提示技术解决了一体化恢复的挑战。具体来说，最先进的PromptIR在解码器端采用隐式提示来学习降级感知提示，旨在完善解码器功能。这种方法虽然很有趣，但并没有扩展到编码器特征的细化，因为编码器特征是未经处理的。尽管隐式提示的想法很有趣，但由于其计算效率低下，仅 37M 参数和 243 GFlops 来处理单个 224×224 大小的图像，因此它对部署提出了相当大的挑战。因此，由于计算要求高，此类模型的实际应用具有挑战性，尤其是在资源有限的手持设备中。然而，选择轻量级模型总是会导致准确性和效率之间的权衡。这些模型中的大多数旨在通过跨步卷积和特征通道分裂来减少参数计数。一些策略也在频域内运行，以减少与注意力机制或非局部操作相关的计算需求，也有的方法通过划分特征空间以实现高效处理，或者将注意力分散到维度上以提高计算效率。

为了在不牺牲性能的情况下优化一体化图像修复效率，本文介绍了一种新颖的权重共享机制。在此方案中，网络模块的权重与其系列的后续模块共享。这种方法大大减少了参数的数量，从而实现了更精简的网络架构。我们将建议的权重共享机制部署在名为动态网络（DyNet）的编码器-解码器风格架构中。除了计算效率外，所提出的DyNet还具有出色的灵活性，因为只需改变模块权重的重用频率，就可以轻松调整网络深度，并在训练期间相应地在笨重和轻量级变体之间切换。

虽然 DyNet 在加速顺序下提供了出色的适应性，但轻量级网络的一个常见挑战是整体准确性的潜在妥协。为了解决这个问题，本文证明了大规模的预训练策略可以有效地提高轻量级模型的性能。通过使用从预训练中得出的权重初始化网络，模型受益于坚实的基础，即使参数数量减少，也可以提高性能。然而，大规模预训练是计算密集型的，需要大量的 GPU 训练时间。因此，本文引入了一种高效且有效的动态预训练策略，能够在单个预训练会话中训练体积更大和更轻量级的网络变体，从而显著节省 50% 的 GPU 训练时间。作为补充，本文编制了一个全面、高质量、高分辨率的数据集，称为 Million-IRD，由 200 万个图像样本组成。本文的主要贡献如下：

• 本文提出了 DyNet，这是一个用于多合一图像恢复任务的动态网络系列。DyNet 可在笨重和轻便之间轻松切换变种。这种灵活性是通过权重共享策略实现的，这是本文网络设计的核心思想，允许有效地重用初始化的模块权重。
• 本文引入了一种动态预训练策略，这是一种新方法，允许在单个会话中同时对体积较大和重量较轻的网络变体进行大规模预训练。这种创新策略将 GPU 训练时间显著减少了 50%，解决了大规模资源密集型模型预训练的紧迫挑战。
• 本文策划了一个全面的预训练数据集，其中包含 200 万张高质量、高分辨率、经过精心过滤的图像。从这个数据集中，本文提取了800万个不重叠的高分辨率补丁，每个补丁的大小为512×512，用于对拟议的DyNet变体进行预训练。

通过这些技术的协同作用，本文提出的DyNet在一体化设置下实现了0.82 dB的图像去噪、去雨和去雾的平均增益，与基线相比，GFlops降低了31.34%，网络参数降低了56.75%;参考下图。

图 1 一体化恢复中的平均 PSNR 与 GFlops 和参数

方法

现有的最先进的一体化图像恢复 （IR） 方法会产生很高的计算占用空间，并且无法在训练期间提供支持具有可变深度的模型的灵活性。为了实现更准确、更轻量级和灵活的一体化IR架构设计，本文提出了一种动态网络架构和一种动态预训练策略。

动态网络 （DyNet） 架构

本文的轻量级架构基于一种用于图像修复模型的新型权重共享机制。这种机制允许网络模块的权重在后续模块之间按顺序重复使用，从而显着减少参数总数并导致更高效的网络结构。本文在编码器-解码器风格的架构中实现了这种权重共享方法，本文称之为动态网络（DyNet）。在 DyNet 中，在每个编码器-解码器级别，模块权重在预先指定数量的后续模块之间共享。这不仅提高了计算效率，还为架构提供了极大的灵活性。通过简单地改变重量分摊的频率，用户可以轻松修改网络的深度，在DyNet的笨重或轻量级变体之间无缝转换。

整体框架

图 2 DyNet的整体框架

所提出的DyNet的整体流水线如上图所示。所提出的 DyNet 首先通过卷积运算从给定的降级输入图像

中提取低级特征

。随后，特征嵌入经过 4 级分层编码器-解码器，逐渐转换为深层潜在特征

。本文利用现有的 Restormer 的 Transformer 模块作为基本特征提取模块。对于编码器和解码器的每个级别，本文初始化第一个Transformer模块的权重，随后在后续模块中以串联的方式重用这些权重，直至该级别的指定频率。例如，在给定的编码器/解码器级别，

表示第一个Transformer模块，我们可以定义该编码器/解码器层的输出

，给定输入特征

为，

其中，

表示自循环。f 表示给定编码器或解码器级别的模块权重重用频率。

本文逐渐将Transformer块的复用频率从顶层提高到底层，从而增加网络深度。多级编码器系统地降低空间分辨率，同时增加通道容量，便于从高分辨率输入图像中逐步提取低分辨率潜在特征。随后，多级解码器逐渐恢复这些低分辨率潜在特征的高分辨率输出。为了增强解码过程，本文集成了提示块，它通过提示生成模块（PGM）和提示交互模块（PIM）隐式学习退化感知提示。与现有的 PromptIR 不同，本文的方法在跳过连接（skip connection）时结合了可识别降级的隐式提示块。这种通过跳过连接的隐式提示可以将退化感知编码器功能传输到解码器端，从而增强恢复过程。网络设计中的这种基本修正为本文一体化图像恢复任务提供了显著改进。

结构细节

在每个编码器-解码器级别，只需更改模块权重重用频率 （f），本文就可以获得 DyNet 的笨重和轻量级变体。在每个编码器级别，本文初始化第一个转换器模块的权重（

），并将其重用于后续模块。本文改变了每个编码器-解码器级别的重用频率，并获得了更笨重和更轻便的变体。DyNet 由四个编码器-解码器级别组成，大型变体 （DyNet-L） 分别在编码器-解码器级别 1 到 4 级别采用

的重用频率。相比之下，轻量级变体 （DyNet-S） 将编码器-解码器级别 1 分别应用于权重重用频率

。此外，本文在每个跳过连接处都加入了一个提示块，这增强了从编码器传输到解码器的编码器特征。总体而言，每个级别都部署了三个提示块，每个提示块由五个提示组件组成。

动态预训练策略

近年来，大规模预训练已成为提高给定网络性能的关键策略。使用预训练权重初始化网络提供坚实的基础，即使使用更少的参数也能提高性能。然而，这种策略是资源密集型的，需要相当大的计算能力和 GPU 训练时间。因此，我们提出了一种动态预训练，能够同时训练网络的多个变体。这些网络由共享网络模块统一，重量不同，但网络深度不同。该策略支持针对各种计算需求和任务复杂性量身定制的各种模型集的并发训练，同时利用共享基线架构。通过建议的动态预训练策略，我们在一次训练中同时训练 DyNet-L 和 DyNet-S 版本的 DyNet，从而减少了 50% 的 GPU 训练时间。

DyNet-L 和 DyNet-S 对在每个编码器-解码器级别初始化的转换器模块使用相同的基础权重。它们只是在每个编码器-解码器级别上这些初始化的转换器模块的重用频率不同。因此，在训练迭代期间，我们在 DyNet-L 和 DyNet-S 之间随机交替，确保优化共享的基础权重，如下图所示。此外，我们通过采用类似于掩蔽自动编码器提出的输入掩蔽策略，提高了DyNet变体的泛化能力。我们随机遮罩图像的某些部分并训练 DyNet 以自我监督的方式重建这些掩蔽区域的变体。用于此训练的数据集如下所述。

图 3 动态预训练策略

Million-IRD：图像修复数据集

图像恢复的大规模预训练实际上需要大规模、高质量和高分辨率的数据集。当前的图像恢复数据集合并后仅提供几千张图像。与用于其他高级任务（如视觉识别、对象检测和分割）的预训练中使用的广泛大规模数据集相比，这远远不够。用于映像恢复的现有训练集的规模相对较小，这限制了底层网络的性能能力。更重要的是，我们受到高级任务中缩放定律的激励，这些定律表明，大规模预保留甚至可以使轻量级、高效的模型设计达到更重模型的性能标记。为了解决这一差距，本文引入了一个名为 Million-IRD 的新百万级数据集，该数据集具有 ∼2M 高质量、高分辨率图像，专门用于预训练图像恢复任务的模型。下图显示了 Million-IRD 数据集中的示例。我们的数据收集和预处理管道将在下面讨论。

图 4 Million-IRD数据集

数据后处理

总体而言，本文的 Million-IRD 数据集拥有 209 万张高质量、高分辨率的图像。我们数据集中的每张图像都带有元数据，详细说明了其下载链接和图像分辨率。补充材料中提供了图像及其各自来源的详细分类。我们从每张图像中提取高分辨率的非重叠斑块（空间大小为512x512），然后应用平坦区域检测器，消除任何包含超过50%平坦面积的斑块。这个后处理阶段使我们能够组装一个由 800 万个不同的图像补丁组成的池，为我们的 DyNet 变体的预训练阶段量身定制。

实验

本文在三个关键的图像恢复任务中验证了所提出的 DyNet：去雾、去雨和去噪。与现有的PromptIR一致，本文的实验在两种不同的设置下进行：（a）多合一，其中单个模型被训练以处理所有三种退化类型，以及（b）单任务，其中单个模型被训练为每个特定的图像恢复任务。

实验细节

动态预训练

为了稳健的权重初始化，本文对 DyNet 的两个变体（即 DyNet-L 和 DyNet-S）进行了动态预训练。这两个变体具有相同的权重，但在每个编码器-解码器级别上的变压器块重用频率不同。对于编码器-解码器1到4级，本文设置变压器块权重复用频率为

( DyNet-L )和

( DyNet-S )。本文使用Million - IRD数据集的全部

大小的8M块进行动态预训练。我们从每个斑块中随机裁剪

个区域，批大小为32。使用JPEG压缩、高斯噪声或不同程度的随机噪声，这些裁剪后的图像块中的每一个都被随机处理其面积的50 %。此外，本文屏蔽了30 %的补丁，而剩下的部分保持不变。总的来说，80 %的输入补丁经过修改，或者通过退化，或者通过掩蔽。对于权重优化，我们使用参数

和

的L1损失和Adam优化器，设置学习率为1e - 4，迭代次数为100万次。使用上述设置，本文同时预训练DyNet - L和DyNet - S。在任意给定的迭代中，本文随机地在两种变体之间切换。有趣的是，随着每次迭代，共享的底层权重被优化。因此，在单次预训练结束时，我们得到的Dy Net - L和Dy Net - S共享相同的训练底层权重，但网络深度不同，使其适合各种挑战，包括鲁棒性和效率。

数据集

对于一体化和单任务实现，我们遵循与现有PromptIR相同的训练协议，并对预训练的DyNet-S和DyNet-L进行微调。我们为各种恢复任务准备数据集。具体来说，对于图像去噪，我们结合BSD400( 400张训练图像)和WED( 4744幅图像)数据集，在

水平上添加高斯噪声来创建含噪图像。在BSD68和Urban100数据集上进行测试。为了去雨，我们使用Rain100L，训练了200张图片和100张测试图片。去雾采用SOTS数据集，包含72135张训练图像和500张测试图像。为了为所有任务开发一个统一的模型，我们将这些数据集合并，微调预训练的DyNet-L和DyNet-S共120轮。

All-in-one设置下多重退化的比较

表 1 All-in-one修复的比较结果

我们将DyNet与一系列通用的图像复原模型和特定的一体化解决方案进行了比较，如上表所示。平均而言，在不同的任务中，DyNet-L和DyNet-S比之前最好的PromptIR分别提高了0.82dB和0.59dB，Dy Net-S也减少了56.75 %的参数和31.34 %的GFlops。在图像去噪方面，与DL相比，DyNet-S的平均PSNR提高了1.09dB。值得注意的是，DyNet-L在图像去雨和去雾方面树立了新的标杆，PSNR分别提高了2.34 d B和1.4 d B。

单一退化的比较

表 2 去雾结果

表 3 去雨结果

表 4 去噪结果

本文评估了DyNet-L在单任务环境下的有效性，其中不同的模型针对特定的修复任务进行了裁剪，证明了通过提示块进行内容自适应提示的有效性。本文的结果，如上表所示表明DyNet-L在去雾方面的性能优越，与PromptIR相比提高了0.76 dB，与Restormer相比提高了1.2 dB。这种模式在其他任务中是一致的，包括去雨和去噪。值得注意的是，与PromptIR相比，DyNet-L在Urban100数据集上取得了1.81 dB的去噪性能增益，在σ = 50的噪声水平上取得了0.13 d B的去噪性能提升。

消融实验

表 5 消融实验

为了研究不同组件的影响，本文在一体化设置下进行了各种消融实验，总结如上表所示。在Million-IRD数据集上，在没有提出动态预训练的情况下，DyNet-L和DyNet-S取得了与基线PromptIR相当的性能，在参数上减少了56.75 %，在GFlops上减少了31.34 %。这表明，与PromptIR中给出的在解码器端放置提示块相比，本文提出的在跳跃连接处放置提示块的基本校正是有效的。此外，在我们的Million - IRD数据集上进行预训练带来了显著的提升，与未进行预训练的DyNet-L相比，DyNet-L的PSNR平均提高了0.55 dB。DyNet-S 变体也观察到类似的效果。此外，本文提出的动态训练策略在训练 DyNet 变体时将 GPU 预训练时间减少了 50%。总体而言，本文的每一项贡献（即通过跳过连接的隐式提示、动态网络架构、Million-IRD 数据集和动态预训练策略）都协同提高了所提方法的性能，同时与基线 PromptIR 相比显着降低了 GFlops 和学习参数。

具有不同退化组合的训练模型

表 6 不同退化组合

本文在聚合数据集上比较了 DyNet-S 在多合一设置中的性能，评估了不同的退化类型组合如何影响其有效性。本消融实验中的所有模型都经过了 80 个轮次的训练。本文评估了退化类型（任务）的不同组合如何影响 DyNet 的性能。值得注意的是，与在去雨和去雾相结合的情况下训练的 DyNet 相比，针对去雨和去噪任务组合训练的 DyNet 在图像去雨方面表现出更好的性能。这表明，当以联合方式训练时，某些退化与模型更相关，可以从中受益。

总结

本文在动态网络（DyNet）中引入了一种新的权重共享机制，用于高效的一体化图像恢复任务，通过性能提升显著提高计算效率。通过调整模块重量重用频率，DyNet 允许在笨重和轻量级模型之间无缝切换。所提出的动态预训练策略在一次训练中同时训练体积大和重量轻的模型。因此，与传统训练策略相比，可节省 50% 的 GPU 预训练时间。消融研究表明，在我们的 Million-IRD 数据集上提出的动态大规模预训练中，笨重和轻量级模型的准确性都显着提高。总体而言，与基线模型相比，DyNet 显着提高了一体化图像修复性能，GFlops 降低了 31.34%，参数减少了 56.75%。