Contrastive Learning for Label-Efficient Semantic Segmentation
This work addresses the high cost of data labeling for semantic segmentation, offering a solution for practitioners and researchers working with limited annotation budgets.
This paper tackles the problem of semantic segmentation with limited labeled data, where deep CNNs tend to overfit. The authors propose a two-stage training strategy: pretraining with a pixel-wise, label-based contrastive loss, followed by fine-tuning with cross-entropy loss, achieving over 20% absolute improvement in some settings and matching or outperforming ImageNet pretraining without additional data.
Collecting labeled data for the task of semantic segmentation is expensive and time-consuming, as it requires dense pixel-level annotations. While recent Convolutional Neural Network (CNN) based semantic segmentation approaches have achieved impressive results by using large amounts of labeled training data, their performance drops significantly as the amount of labeled data decreases. This happens because deep CNNs trained with the de facto cross-entropy loss can easily overfit to small amounts of labeled data. To address this issue, we propose a simple and effective contrastive learning-based training strategy in which we first pretrain the network using a pixel-wise, label-based contrastive loss, and then fine-tune it using the cross-entropy loss. This approach increases intra-class compactness and inter-class separability, thereby resulting in a better pixel classifier. We demonstrate the effectiveness of the proposed training strategy using the Cityscapes and PASCAL VOC 2012 segmentation datasets. Our results show that pretraining with the proposed contrastive loss results in large performance gains (more than 20% absolute improvement in some settings) when the amount of labeled data is limited. In many settings, the proposed contrastive pretraining strategy, which does not use any additional data, is able to match or outperform the widely-used ImageNet pretraining strategy that uses more than a million additional labeled images.