Augmentation-Free Dense Contrastive Knowledge Distillation for Efficient Semantic Segmentation
作者: Jiawei Fan, Chao Li, Xiaolong Liu, Meina Song, Anbang Yao
分类: cs.CV, cs.AI, cs.LG
发布日期: 2023-12-07
备注: The paper of Af-DCD is accepted to NeurIPS 2023. Code and models are available at https://github.com/OSVAI/Af-DCD
🔗 代码/项目: GITHUB
💡 一句话要点
提出无数据增强的密集对比知识蒸馏方法，提升语义分割效率与精度。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 语义分割 知识蒸馏 对比学习 无数据增强 模型压缩
📋 核心要点
- 现有语义分割知识蒸馏方法依赖数据增强和内存缓冲，计算资源需求高，限制了其应用。
- Af-DCD通过掩码特征模仿和特征分区对比学习损失，实现高效的密集知识迁移。
- 实验表明，Af-DCD在多个数据集上显著提升了学生模型的mIOU，并创造了新的性能记录。
📝 摘要（中文）
近年来，基于对比学习的知识蒸馏方法在图像分类和目标检测任务上取得了显著成果。然而，在语义分割领域，相关研究较少。现有方法严重依赖数据增强和内存缓冲，这在处理需要保留高分辨率特征图以进行密集像素预测的语义分割任务时，会带来很高的计算资源需求。为了解决这个问题，我们提出了一种无数据增强的密集对比知识蒸馏（Af-DCD）方法，这是一种新的对比蒸馏学习范式，用于训练紧凑而精确的深度神经网络，以用于语义分割应用。Af-DCD利用掩码特征模仿策略，并通过巧妙地利用通道和空间维度上的特征分区，构建了一种新的对比学习损失，从而有效地将教师模型学习到的密集和结构化的局部知识转移到目标学生模型，同时保持训练效率。在五个主流基准数据集上，使用各种教师-学生网络对进行的大量实验证明了我们方法的有效性。例如，当选择DeepLabV3-Res101作为教师模型时，由Af-DCD训练的DeepLabV3-Res18|DeepLabV3-MBV2模型在Cityscapes数据集上分别达到了77.03%|76.38%的mIOU，创造了新的性能记录。此外，与单独训练的对应模型相比，Af-DCD在Cityscapes|Pascal VOC|Camvid|ADE20K|COCO-Stuff-164K上实现了3.26%|3.04%|2.75%|2.30%|1.42%的绝对mIOU提升。
🔬 方法详解
问题定义：现有语义分割知识蒸馏方法为了提升性能，通常依赖于大量的数据增强和维护一个庞大的内存缓冲池。数据增强增加了计算负担，而内存缓冲池则需要大量的存储空间。这些需求使得现有方法在处理高分辨率图像的语义分割任务时效率低下，难以应用到资源受限的场景中。
核心思路：Af-DCD的核心思路是通过一种无数据增强的对比学习方法，将教师模型学习到的密集和结构化的局部知识有效地迁移到学生模型。该方法避免了数据增强带来的额外计算开销，并利用特征分区策略构建对比学习损失，从而实现高效的知识迁移。
技术框架：Af-DCD主要包含两个关键模块：掩码特征模仿和对比学习损失。首先，使用掩码特征模仿策略，让学生模型学习教师模型的特征表示。然后，通过在通道和空间维度上对特征进行分区，构建正负样本对，并设计对比学习损失，促使学生模型学习教师模型中像素级别的关系。整体训练流程为：输入图像，教师模型和学生模型分别提取特征，利用掩码特征模仿策略对齐特征，计算对比学习损失，更新学生模型参数。
关键创新：Af-DCD的关键创新在于提出了无数据增强的密集对比学习框架。与现有方法相比，Af-DCD不需要数据增强，降低了计算复杂度。此外，通过特征分区构建对比学习损失，能够有效地迁移教师模型中的密集和结构化知识，提升学生模型的性能。
关键设计：Af-DCD的关键设计包括：1) 掩码特征模仿策略，通过随机掩盖部分特征，增强学生模型的鲁棒性；2) 特征分区策略，在通道和空间维度上对特征进行划分，构建正负样本对，用于对比学习；3) 对比学习损失函数，促使学生模型学习教师模型中像素级别的关系。具体的损失函数形式未知，需要参考论文原文。
📊 实验亮点
Af-DCD在Cityscapes数据集上，使用DeepLabV3-Res101作为教师模型时，DeepLabV3-Res18和DeepLabV3-MBV2学生模型分别达到了77.03%和76.38%的mIOU，创造了新的性能记录。与单独训练的对应模型相比，Af-DCD在Cityscapes、Pascal VOC、Camvid、ADE20K和COCO-Stuff-164K数据集上分别实现了3.26%、3.04%、2.75%、2.30%和1.42%的绝对mIOU提升。
🎯 应用场景
该研究成果可应用于自动驾驶、遥感图像分析、医疗图像分割等领域。通过知识蒸馏，可以将大型复杂模型中的知识迁移到小型轻量级模型中，从而在资源受限的设备上实现高性能的语义分割，具有重要的实际应用价值和广泛的应用前景。
📄 摘要（原文）
In recent years, knowledge distillation methods based on contrastive learning have achieved promising results on image classification and object detection tasks. However, in this line of research, we note that less attention is paid to semantic segmentation. Existing methods heavily rely on data augmentation and memory buffer, which entail high computational resource demands when applying them to handle semantic segmentation that requires to preserve high-resolution feature maps for making dense pixel-wise predictions. In order to address this problem, we present Augmentation-free Dense Contrastive Knowledge Distillation (Af-DCD), a new contrastive distillation learning paradigm to train compact and accurate deep neural networks for semantic segmentation applications. Af-DCD leverages a masked feature mimicking strategy, and formulates a novel contrastive learning loss via taking advantage of tactful feature partitions across both channel and spatial dimensions, allowing to effectively transfer dense and structured local knowledge learnt by the teacher model to a target student model while maintaining training efficiency. Extensive experiments on five mainstream benchmarks with various teacher-student network pairs demonstrate the effectiveness of our approach. For instance, the DeepLabV3-Res18|DeepLabV3-MBV2 model trained by Af-DCD reaches 77.03%|76.38% mIOU on Cityscapes dataset when choosing DeepLabV3-Res101 as the teacher, setting new performance records. Besides that, Af-DCD achieves an absolute mIOU improvement of 3.26%|3.04%|2.75%|2.30%|1.42% compared with individually trained counterpart on Cityscapes|Pascal VOC|Camvid|ADE20K|COCO-Stuff-164K. Code is available at https://github.com/OSVAI/Af-DCD