以恒定内存将数据集蒸馏扩展至 ImageNet-1K
计算机视觉与模式识别
2023-11-02 v4 人工智能
摘要
数据集蒸馏是一个新兴领域，旨在将大型数据集蒸馏为更小且信息高度密集的合成数据集，以加速训练并减少存储。在各种数据集蒸馏方法中，基于轨迹匹配的方法(MTT)已在诸多任务上取得 SOTA 性能，例如在 CIFAR-10/100 上。然而，由于通过 SGD 步骤展开优化时内存消耗过高，MTT 无法扩展至 ImageNet-1K 等大规模数据集。我们能否将该 SOTA 方法扩展至 ImageNet-1K，且其在 CIFAR 上的有效性是否能迁移到 ImageNet-1K？为回答这些问题，我们首先提出一种以恒定内存复杂度精确计算展开梯度的方法，使我们能将 MTT 无缝扩展至 ImageNet-1K，并将内存占用降低约 6 倍。我们进一步发现 MTT 难以处理具有大量类别的数据集，并提出一种新颖的软标签分配方法，大幅改善其收敛性。所得算法在 ImageNet-1K 上刷新 SOTA：我们能在单块 GPU 上将 IPCs(每类图像数)扩展至 50(此前所有方法在 ImageNet-1K 上仅能扩展至 2 IPCs)，在仅使用 4.2% 数据量的同时取得最佳精度(相较全数据集训练仅下降 5.9% 精度)——较先前 SOTA 绝对提升 18.2%。我们的代码见 https://github.com/justincui03/tesla
引用
@article{arxiv.2211.10586,
title = {Scaling Up Dataset Distillation to ImageNet-1K with Constant Memory},
author = {Justin Cui and Ruochen Wang and Si Si and Cho-Jui Hsieh},
journal= {arXiv preprint arXiv:2211.10586},
year = {2023}
}