Confidence-Aware Contrastive Learning for Semantic Segmentation

摘要

Recently supervised contrastive learning (SCL) has achieved remarkable progress in semantic segmentation. Nevertheless, prior works have often necessitated a substantial number of samples to attain satisfactory performance, leading to a significant increase in training overhead. In this work, we leverage the idea of reweighting each pair to reduce the demand for large numbers of training samples in contrastive learning and propose a novel loss, dubbed confidence-aware contrastive (CAC) loss, which adaptively reweights each pair according to the predicted confidence for semantic segmentation. To alleviate the misalignment between supervised learning and contrastive learning, we further introduce an extra weight branch with a stop-gradient operator to generate the pair weights. Moreover, we present a confidence-aware marginal anchor sampling method for the calculation of supervised contrastive loss which focuses on marginal rather than the hardest pairs. Coupling with our method consistently improves the performance of various models (e.g. HRNet, OCRNet, SegFormer) on Cityscapes, ADE20K, PASCAL-Context, and COCO-Stuff datasets. Compared to existing SCL-based methods, the proposed method achieves competitive or even better results without relying on a memory bank or a large number of samples. Our code is at https://github.com/CVIU-CSU/Confidence-Aware-Contrastive-Loss.