中文

ReDistill:面向卷积神经网络峰值内存缩减的残差编码蒸馏

计算机视觉与模式识别 2025-04-29 v3 机器学习

摘要

神经网络规模的扩大和现代图像传感器分辨率的提升,导致处理现代计算机视觉模型所需的内存和功耗需求增加。为了在资源极度受限的边缘设备上部署这些模型,减少其峰值内存(即模型执行期间消耗的最大内存)至关重要。一种降低峰值内存的朴素方法是通过大步长池化对特征图进行激进下采样,但这通常会导致网络性能不可接受的下降。为缓解此问题,我们提出了用于在教师-学生框架中减少峰值内存的残差编码蒸馏(ReDistill),其中通过使用激进池化从教师网络派生出一个内存更少的学生网络。我们将我们的蒸馏方法应用于计算机视觉中的多个问题,包括图像分类和基于扩散模型的图像生成。对于图像分类,我们的方法在大多数基于CNN的架构上实现了4-5倍的理论峰值内存缩减,且精度下降更少。对于基于扩散模型的图像生成,我们提出的蒸馏方法生成了一个理论峰值内存降低4倍的去噪网络,同时保持了良好的图像生成多样性和保真度。实验表明,当应用于相同的学生网络时,我们的方法优于其他基于特征和基于响应的蒸馏方法。代码可在 https://github.com/mengtang-lab/ReDistill 获取。

关键词

引用

@article{arxiv.2406.03744,
  title  = {ReDistill: Residual Encoded Distillation for Peak Memory Reduction of CNNs},
  author = {Fang Chen and Gourav Datta and Mujahid Al Rafi and Hyeran Jeon and Meng Tang},
  journal= {arXiv preprint arXiv:2406.03744},
  year   = {2025}
}

备注

16 pages, 7 figures, 10 tables