[All-in-one] 动态预训练 Dynamic Pre-training : 面向高效和可扩展的一体化图像恢复-CSDN博客

本文链接：https://blog.csdn.net/u014546828/article/details/141302480

Dynamic Pre-training: Towards Efficient and Scalable All-in-One Image Restoration

Abstract

All-in-one image restoration tackles different types of degradations with a unified model instead of having task-specific, non-generic models for each degradation. The requirement to tackle multiple degradations using the same model can lead to high-complexity designs with fixed configuration that lack the adaptability to more efficient alternatives. We propose DyNet, a dynamic family of networks designed in an encoder-decoder style for all-in-one image restoration tasks. Our DyNet can seamlessly switch between its bulkier and lightweight variants, thereby offering flexibility for efficient model deployment with a single round of training. This seamless switching is enabled by our weightssharing mechanism, forming the core of our architecture and facilitating the reuse of initialized module weights. Further, to establish robust weights initialization, we introduce a dynamic pre-training strategy that trains variants of the proposed DyNet concurrently, thereby achieving a 50% reduction in GPU hours. To tackle the unavailability of largescale dataset required in pre-training, we curate a high-quality, highresolution image dataset named Million-IRD having 2M image samples. We validate our DyNet for image denoising, deraining, and dehazing in all-in-one setting, achieving state-of-the-art results with 31.34% reduction in GFlops and a 56.75% reduction in parameters compared to baseline models.

动态预训练：迈向高效且可扩展的一站式图像恢复

一站式图像恢复通过统一的模型来处理不同类型的退化问题，而不是为每个退化问题设计任务特定的非通用模型。使用同一模型处理多种退化的需求可能导致设计复杂度增加，且固定配置可能缺乏适应更高效替代方案的能力。

本文提出了DyNet，一个为一体化图像恢复任务设计的动态网络家族，采用编码器-解码器风格。DyNet可以在其更庞大和轻量级的变体之间无缝切换，从而在单次训练后提供灵活高效的模型部署方案。这种无缝切换是通过权重共享机制实现的，这是本文架构的核心，并促进了初始化模块权重的重用。

此外，为了建立稳健的权重初始化，本文引入了一种动态预训练策略，该策略同时训练DyNet的多个变体，从而将GPU时间减少了50%。为了解决预训练所需的大规模数据集不可用的问题，我们整理了一个名为Million-IRD的高质量、高分辨率图像数据集，包含200万张图像样本。

本文验证了一站式设置下DyNet在图像去噪、去雨和去雾方面的性能，与基线模型相比，本文的方法在GFLOPs上减少了31.34%，在参数上减少了56.75%，同时达到了最先进的结果。这证明了DyNet在高效性和可扩展性方面的一站式图像恢复能力。

Introduction

图像恢复任务的挑战

多样性和严重性：图像中存在的多种类型和不同程度的退化（如噪声、雨痕、雾等）对图像恢复提出了重大挑战。
现有方法的局限性：大多数现有方法通过隐式学习图像先验来恢复图像，这要求针对不同的退化类型、程度和数据集进行单独的网络训练。这些方法在测试时还需要关于图像退化的先验知识来选择有效的模型，因此缺乏处理多种退化的通用性。

All-in-One图像恢复的目标

统一模型：All-in-One图像恢复旨在使用单个统一模型来恢复具有未知退化的图像，而不是为每个退化类型训练专门的、非通用的模型。
现有方法的不足：尽管最近的进展如AirNet和PromptIR通过对比学习和隐式视觉提示技术解决了All-in-One恢复挑战，但这些方法仍存在计算效率低下的问题，特别是在处理高分辨率图像时。例如，PromptIR虽然有趣，但由于其高达37M的参数和243GFlops的计算量，在实际应用中面临挑战。

本文方法

DyNet网络：本文提出了一种名为DyNet的动态网络家族，采用编码器-解码器风格，旨在解决All-in-One图像恢复任务。DyNet能够在其庞大和轻量级变体之间无缝切换，从而在单次训练中提供灵活的模型部署。
权重共享机制：DyNet的核心是其权重共享机制，该机制允许网络模块的权重在其后续模块中重用，从而显著减少参数数量并优化网络结构。
动态预训练策略：为了建立鲁棒的权重初始化，本文引入了一种动态预训练策略，该策略能够同时训练DyNet的庞大和轻量级变体，从而在单个预训练会话中节省50%的GPU时间。
Million-IRD数据集：为了应对大规模预训练所需数据集不可用的挑战，本文整理了一个包含200万图像样本的高质量、高分辨率图像数据集Million-IRD。

Method

1. Dynamic Network (DyNet) Architecture

权重共享机制：

核心思想：DyNet采用权重共享机制作为其设计的核心，这一机制允许在后续的模块中高效地重用已初始化的模块权重，从而显著减少总参数数量并提升网络效率。
实现方式：权重共享在编码器-解码器风格的架构中实施，使得网络模块的权重能够在序列中的后续模块中重复使用。

网络结构设计：

编码器-解码器架构：DyNet设计为一个编码器-解码器风格的架构，旨在处理所有类型的图像恢复任务。
模块配置：在每个编码器-解码器级别，通过调整模块权重共享的频率，可以轻松地在更庞大和轻量级的模型变体之间切换。
Transformer块：网络中包含Transformer块，每个块由两个子模块组成：多Dconv头转置注意力模块（MDTA）和门控Dconv前馈网络（GDFN）。MDTA模块专注于计算跨通道维度的注意力，而非空间维度，这显著降低了计算需求。

灵活性：

深度调整：通过简单地改变模块权重重用的频率，可以轻松调整网络的深度，从而在庞大和轻量级的模型之间无缝切换。
任务适应性：DyNet能够在单次训练后灵活部署于不同的计算环境中，以满足不同的性能和资源需求。

特征一致性：

跳跃连接：DyNet通过跳跃连接隐式地学习退化感知提示，以保持编码器-解码器特征的一致性，这与在解码器侧进行学习的PromptIR方法不同。
效果：这种设计使得DyNet能够更有效地恢复图像质量，同时保持特征的稳定性和一致性。

技术细节：

权重重用：网络中的每个模块通过共享底层权重实现权重重用，这种设计显著减少了模型的总参数量。
参数效率：由于权重共享，DyNet在保持高性能的同时，实现了显著的参数减少和计算效率提升。

2. Dynamic Pre-training Strategy

近年来，大规模预训练已成为提升给定网络性能的关键策略。使用预训练权重初始化网络为模型提供了一个坚实的基础，即使参数较少也能显著提升性能。然而，这一策略资源密集，需要大量的计算能力和GPU时间。

因此，本文提出了一种动态预训练策略，能够同时训练网络的多个变体。这些网络通过共享网络模块的权重实现统一，但网络深度有所不同。这一策略允许同时训练一系列针对不同计算需求和任务复杂度的定制模型，同时利用共享的基线架构。通过提出的动态预训练策略，在单次训练会话中同时训练了DyNet的DyNet-L和DyNet-S版本，实现了GPU时间减半的效果。

DyNet-L和DyNet-S在每个编码器-解码器级别初始化的Transformer块使用相同的基础权重。它们之间的主要区别在于每个编码器-解码器级别中这些初始化Transformer块的复用频率。因此，在训练迭代过程中，随机在DyNet-L和DyNet-S之间切换，以确保如图3所示优化共享的基础权重。此外，我们采用了一种类似于掩码自编码器[19]中提出的输入掩码策略，以增强DyNet变体的泛化能力。我们随机掩蔽图像的部分区域，并训练DyNet变体以自监督的方式重建这些掩蔽区域。用于此训练的数据集将在下文描述。