用智能选图+分阶段训练,让修复图像仅用1张消费级显卡8小时搞定
Distribution-aware Dataset Distillation for Efficient Image Restoration
- 按图像复杂度分阶段挑选样本,先练简单图再攻难题
- 用轻量CNN调整特征分布,使小数据集逼近原数据效果
- 4K超高清图像修复只需1张显卡8小时,省下500倍算力
面对图像数据爆炸式增长,训练图像修复模型成本高昂。现有数据蒸馏技术在图像修复领域几乎空白。为此,我们提出分布感知的数据蒸馏方法TripleD,将数据蒸馏扩展至图像修复任务。TripleD利用预训练视觉Transformer提取图像特征并评估复杂度,基于复杂度筛选出远小于原始数据集的子集。该子集经轻量CNN微调,使特征分布与原始数据集对齐。训练分为两阶段:初期聚焦低复杂度样本建立基础能力,后期逐步引入高复杂度和不确定样本。实验表明,TripleD在多任务图像修复、全功能修复及超高清修复任务中表现优异。仅需1张消费级显卡,在不到8小时内即可训练出先进水平的4K图像修复模型,计算资源节省达500倍,训练时间显著缩短。
原文摘要 · Abstract (English)
With the exponential increase in image data, training an image restoration model is laborious. Dataset distillation is a potential solution to this problem, yet current distillation techniques are a blank canvas in the field of image restoration. To fill this gap, we propose the Distribution-aware Dataset Distillation method (TripleD), a new framework that extends the principles of dataset distillation to image restoration. Specifically, TripleD uses a pre-trained vision Transformer to extract features from images for complexity evaluation, and the subset (the number of samples is much smaller than the original training set) is selected based on complexity. The selected subset is then fed through a lightweight CNN that fine-tunes the image distribution to align with the distribution of the original dataset at the feature level. To efficiently condense knowledge, the training is divided into two stages. Early stages focus on simpler, low-complexity samples to build foundational knowledge, while later stages select more complex and uncertain samples as the model matures. Our method achieves promising performance on multiple image restoration tasks, including multi-task image restoration, all-in-one image restoration, and ultra-high-definition image restoration tasks. Note that we can train a state-of-the-art image restoration model on an ultra-high-definition (4K resolution) dataset using only one consumer-grade GPU in less than 8 hours (500 savings in computing resources and immeasurable training time).
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。