Dynamic Pre-training: Towards Efficient and Scalable All-in-One Image Restoration
https://arxiv.org/pdf/2404.02154
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变体以自监督的方式重建这些掩蔽区域。用于此训练的数据集将在下文描述。
Million-IRD数据集
数据集的重要性:
数据集需求:大型、高质量的图像恢复数据集对于有效的预训练至关重要。现有的图像恢复数据集规模较小,限制了预训练模型的性能。
规模对比:与用于高级任务(如视觉识别、目标检测和分割)的广泛使用的大型数据集相比,现有图像恢复数据集的数量明显不足。
Million-IRD数据集的介绍:
数据集规模:Million-IRD是一个包含约200万张高质量、高分辨率图像的数据集,专为图像恢复任务的预训练而构建。
数据集来源:数据集通过从不同源(如NTIRE[4]、DIV2K[2]、Flickr2K[35]、LSDIR[27]和Laion-HR[42])收集和整理图像来构建。
数据集特点:数据集具有高度的多样性和丰富性,涵盖了多种不同的场景和图像内容。
数据集的构建过程:
数据收集:从多个源收集高质量的图像,并进行严格的筛选和预处理。
质量控制:拒绝不符合质量标准的图像,确保数据集的纯净度和高质量。
统计数据:提供了数据集中每个来源的图像数量的详细统计信息(如Table S1所示)。
数据集的应用价值:
预训练优势:通过在大规模数据集上进行预训练,可以显著提高图像恢复模型的性能,即使对于轻量级模型也能达到与更重模型相当的效果。
研究贡献:Million-IRD数据集的引入填补了图像恢复任务中大规模数据集的空白,为未来的研究提供了坚实的基础。
数据集的样例展示:
样本图像:文档展示了Million-IRD数据集中的一些样本图像,以及被拒绝的低质量图像示例(如图4所示)。
数据集的潜在影响:
研究方向推动:Million-IRD数据集的发布将促进图像恢复领域的研究,推动更高效、可扩展的模型的发展。
实际应用潜力:高质量的预训练模型将能够更好地应对实际应用中的图像恢复需求,提高图像处理的效率和质量。
数据集的贡献:
数据集的发布不仅为图像恢复任务的预训练提供了新的资源,还通过展示其在实际应用中的潜力,强调了大数据在推动深度学习模型进步中的重要性。