中文

动态预训练:迈向高效且可扩展的一体化图像复原

计算机视觉与模式识别 2024-10-15 v2

摘要

一体化图像复原使用统一的模型处理不同类型的退化,而非为每种退化配备任务特定的非通用模型。使用同一模型处理多种退化的需求可能导致具有固定配置的高复杂度设计,缺乏对更高效替代方案的适应性。我们提出了 DyNet,这是一种以编码器-解码器风格设计的动态网络家族,用于一体化图像复原任务。我们的 DyNet 可以在其大型与轻量级变体之间无缝切换,从而仅需一轮训练即可为高效的模型部署提供灵活性。这种无缝切换由我们的权重共享机制实现,该机制构成了我们架构的核心并促进了已初始化模块权重的重用。此外,为了建立稳健的权重初始化,我们引入了一种动态预训练策略,该策略并发训练所提 DyNet 的变体,从而将 GPU 小时数减少了 50%。我们的动态预训练策略消除了为每个变体维护单独检查点的需求,因为所有模型共享一组通用的检查点,仅在模型深度上有所不同。这种高效的策略显著降低了存储开销并增强了适应性。为了解决预训练所需的大规模数据集不可用的问题,我们构建了一个名为 Million-IRD 的高质量、高分辨率图像数据集,包含 200 万个图像样本。我们在一体化设置下验证了 DyNet 在图像去噪、去雨和去雾任务上的表现,取得了 state-of-the-art 的结果,与基线模型相比,GFlops 减少了 31.34%,参数量减少了 56.75%。源代码和训练好的模型可在 https://github.com/akshaydudhane16/DyNet 获取。

关键词

引用

@article{arxiv.2404.02154,
  title  = {Dynamic Pre-training: Towards Efficient and Scalable All-in-One Image Restoration},
  author = {Akshay Dudhane and Omkar Thawakar and Syed Waqas Zamir and Salman Khan and Fahad Shahbaz Khan and Ming-Hsuan Yang},
  journal= {arXiv preprint arXiv:2404.02154},
  year   = {2024}
}

备注

This version includes updates where the DyNet variants now share the same weights during inference as well, eliminating the need to store separate weights and thereby reducing device storage requirements. Additionally, all results have been updated based on the new experimental setup