Region-level Contrastive and Consistency Learning for Semi-Supervised Semantic Segmentation

1 Introduction

Semantic segmentation has high potential values in a variety of applications. However, training supervised semantic segmentation models requires large-scale pixel-level annotations, and such pixel-wise labeling is time-consuming and expensive. This work focuses on semi-supervised semantic segmentation, which takes advantage of a large amount of unlabeled data and limits the need for labeled examples.

Currently, many works have demonstrated that designing pixel-level consistency regularization [French et al., 2020; Ouali et al., 2020; Chen et al., 2021], and pixel-level contrastive regularization [Zhong et al., 2021; Lai et al., 2021]

∗Equal contribution

† Interns at the Institute of Deep Learning, Baidu Research

‡Corresponding author

are beneficial for semi-supervised semantic segmentation. The first family of works utilized pixel-level label consistency under different data augmentations [French et al., 2020; Zou et al., 2021], feature perturbations [Ouali et al., 2020], network branches [Ke et al., 2020], and segmentation models [Chen et al., 2021]. These methods benefited from using the label-space consistency property on the unlabeled images. Figure 1 (a) shows the pixel-level label consistency of employing different data augmentations. Differently, the work [Lai et al., 2021] proposed the Directional Contrastive Loss to accomplish the pixel-level feature contrast, which is shown in Figure 1 (b). Substantial progress has been made by utilizing pixel-level label consistency and pixel-level feature contrastive property. For example, PC2Seg [Zhong et al., 2021] leveraged and simultaneously enforced the consistency property in the label space and the contrastive property in the feature space, which is illustrated in Figure 1 (c). However, such pixel-level regularization is sensitive to noise from pixels with incorrect predictions. Besides, pixel-level contrastive regularization has memory and computational cost with O(pixel num2 ), and usually requires designing additional negative example filtering mechanism carefully.

To overcome the above challenges, we propose to design Region-level Contrastive and Consistency Learning (RC2L) for semi-supervised semantic segmentation. As shown in Figure 1 (d), our method enforces the consistency of region classes and the contrastive property of features and masks from different regions. It was inspired by MaskFormer [Cheng et al., 2021] which formulated supervised semantic segmentation as a mask (or region) classification problem.

Specifically, we first propose a Region Mask Contrastive (RMC) loss and a Region Feature Contrastive (RFC) loss to achieve region-level contrastive property. The former pulls the masks of matched region pairs (or positive pairs) closer and pushes away the unmatched region pairs (or negative pairs), and the latter pulls the features of matched region pairs closer and pushes away the unmatched regions pairs. Furthermore, Region Class Consistency (RCC) loss and Semantic Mask Consistency (SMC) loss are proposed for encouraging the consistency of the region classes and the consistency of union regions with the same class, respectively.

Based on the proposed components, we develop a Regionlevel Contrastive and Consistency Learning (RC2L) framework for semi-supervised semantic segmentation. The effectiveness of our approach is demonstrated by conducting extensive ablation studies. In addition, we evaluate our RC2L on several widely-used benchmarks, e.g., PASCAL VOC 2012 and Cityscapes, and the experiment results show that our approach outperforms the state-of-the-art semisupervised segmentation methods.

2 RelatedWorks

Semantic segmentation. Since the emergence of Fully Convolutional Network (FCN) [Long et al., 2015], per-pixel classifications have achieved high accuracies in semantic segmentation tasks. The development of modern deep learning methods for semantic segmentation mainly focuses on how to model context [Chen et al., 2018; Wu et al., 2020b; Wu et al., 2020a]. Recently, Transformer-based methods for semantic segmentation show an excellent performance. SegFormer [Xie et al., 2021] proposed mask transformer as a decoder. More recently, MaskFormer [Cheng et al., 2021] reformulated semantic segmentation as a mask (or region) classification task. Differently, we focus on how to make a better use of unlabeled data to perform region-level predictions for semi-supervised semantic segmentation.

Semi-supervised semantic segmentation. It is important to explore semi-supervised semantic segmentation to reduce per-pixel labeling costs. Early approach [Hung et al., 2018]

used generative adversarial networks (GANs) and adversarial loss to train on unlabeled data. Consistency regularization and pseudo-labeling have also been widely explored for semi-supervised semantic segmentation, by enforcing consistency among the predictions, either from feature perturbation [Ouali et al., 2020], augmented input images [French et al., 2020], or different segmentation models [Ke et al., 2020]. Later, PseudoSeg [Zou et al., 2021] combined pixellevel labels with image-level labels to enhance the semisupervised learning through boosting the quality of pseudo labels. CPS [Chen et al., 2021] proposed the cross pesudo supervision to apply consistency between two segmentation networks with the same architecture. Self-training has also driven the development of the state-of-the-art methods [Zoph et al., 2020; He et al., 2021]. Pixel-level contrastive based approaches have been proposed recently. CMB [Alonso et al., 2021] introduced a memory bank for positive-only contrastive learning. Directional Contrastive Loss [Lai et al., 2021] was proposed for training between pixel features and pseudo labels. PC2Seg [Zhong et al., 2021] enforced the pseudo labels to be consistent, and encouraged pixel-level contrast to pixel features. All these methods mainly focused on pixel-level regularization and conducted pixel-level predictions. Differently, our approach explores region-level regularization and performs region-level predictions.

Contrastive learning. Image-level contrastive learning has shown excellent prospects for self-supervised representation learning. SimCLR [Chen et al., 2020a] proposed to take the contrastive loss for training between images by applying different data augmentations. MoCo V2 [Chen et al., 2020b] presented a momentum encoder to reduce the requirement of large batch size. Pixel-level contrastive learning has also been proven to be beneficial for dense prediction tasks. [Wang et al., 2021] introduced pixel-level contrastive learning for supervised semantic segmentation. Different from these works, we explore how to conduct region-level contrastive learning for semi-supervised semantic segmentation.

3 Method

We first describe our framework, Region-level Contrastive and Consistency Learning (RC2L) for semi-supervised semantic segmentation. Then, we present the Region Mask Contrastive (RMC) loss and Region Feature Contrastive (RFC) loss. The former pulls the masks of matched region pairs closer and pushes away the unmatched region pairs, and the latter pulls the features of matched region pairs closer and pushes away the unmatched region pairs. Finally, we introduce the Region Class Consistency (RCC) loss and Semantic Mask Consistency (SMC) loss. RCC loss can enforce the consistency of region classes, and SMC loss can further promote the consistency of union regions with the same class.

3.1 Framework

The overall framework of our RC2L is shown in Figure 2, which consists of a student model and a teacher model. The teacher model has the same architecture as the student model but uses a different set of weights. The student model is trained with both labeled and unlabeled data.

Specifically, given an image-label pair, {xl, zgt }, xl is the image, and zgt is the corresponding annotations that are the set of N gt ground truth segments, i.e., zgt = {(cgt i ,mgt i )|cgt i ∈ {1, 2, ...,K},mgt i ∈ {0, 1}H×W }N gt i=1 , where ci is the ground truth class of the ith segment, W and H represent the width and height of the input image. We feed the image xl into the student model, outputting probability-mask pairs zl = {(pli ,mli )|pli ∈ △K+1 ,mli ∈ [0, 1]H×W }Ni=1 , where the probability distribution pi contains a "no object" label ⊘ and K category labels. A bipartite matching-based assignment σ between the set of predictions zl and zgt is conducted for computing supervised loss:

where Lmask is a binary mask loss. Please refer to [Cheng et al., 2021] for more details.

For the unlabeled image xu, we feed its weak augmentation xwu and strong augmentation xsu into the teacher and student models, respectively. Here, we use short edge resize, random crop, random flip, and color augmentation as weak augmentations, and use all weak data augmentation methods and CutMix [Yun et al., 2019] as strong augmentations. The teacher is employed to generate pseudo label zt = {(cti ,mti )}N t i=1 which is used to guide the training of the student. The unsupervised loss can be formulated as follow:

where LRCC and LSMC mean the Region Class Consistency loss and Semantic Mask Consistency loss (Section 3.2), respectively. LRMC and LRFC denote the Region Mask Contrastive loss and Region Feature Contrastive loss (Section 3.3), respectively. β1, β2, β3 and β4 denote the loss weight. Therefore, the total loss is defined as:

where α is a constant to balance between the supervised and unsupervised losses.

Following Mean Teacher [Tarvainen and Valpola, 2017], the parameters θt of teacher model are an exponential moving average of parameters θs of the student model. Specifically, at every training step, the parameters θt of teacher network is updated as follows:θt = τ θt + (1 - τ )θs, (4)

where τ ∈ [0, 1 ] is a decay rate.

3.2 Region-level Contrastive Learning

The goal of region-level contrastive learning is to increase the similarity between each region mask and feature from strong augmentation xsu and the mask and feature of the matched region (from weak augmentation xwu ), and reduce the similarity between unmatched region pairs. Specifically, we propose a Region Mask Contrastive (RMC) loss and a Region Feature Contrastive (RFC) loss to achieve region-level contrastive property. The former pulls the masks of matched region pairs (or positive pairs) closer and pushes away the unmatched region pairs (or negative pairs), and the latter pulls the features of matched region pairs closer and pushes away the unmatched regions pairs.

The weak augmentation xwu is firstly fed into the teacher network, which outputs pseudo labels zt = {(cti ,mti )}N t i=1 . The strong augmentation xsu is fed into the student network, which outputs class logits ps = {psi ∈ △K+1 }Ni=1 , mask embeddings f s = {f s i ∈ RC }Ni=1 , and per-pixel features F s ∈ RC×H/4×W/4 . Firstly, we obtain binary mask predictions {msi }Ni=1 via a dot product between the mask embeddings and per-pixel features. Then, the bipartite matchingbased assignment σ between student predictions {msi }Ni=1 and pseudo segment set {mti }N t i=1 is used for getting matched index set ID = {σ(i)}N t i=1 and computing RMC Loss:

where τm is a temperature hyper-parameter to control the scale of terms inside exponential, and d(msi ,mtj ) = 2|msi ∩mtj | |msi |∪|mtj | measures the similarity between two masks.

Next, we compute region features rs = {rsσ(i) }N t i=1 via combining per-pixel features F s and region mask set {msi }Ni=1 . The process can be formulated as follow:

where "GAP" indicates a global average pooling operation. We compute target region features rt = {rt i }N t i=1 as follow:

Here, per-pixel features F t are not used to compute the target region features, since the student model and the teacher model have different weights and feature space. Then, Region Feature Contrastive loss is defined as:

where τf is a temperature hyper-parameter to control the scale of terms inside exponential, cos(u, v) = uT v ∥u∥∥v∥ is the cosine similarity.

Comparison with Pixel Contrastive Loss. The proposed Region Mask Contrastive (RMC) loss and Region Feature Contrastive (RFC) loss are different from the most related Pixel Contrastive Loss [Zhong et al., 2021] in two aspects. Firstly, it performed contrastive learning on pixel-level features, while our method conducts contrastive learning on region-level masks and features. Secondly, unlike [Zhong et al., 2021], our methods do not require designing the negative example filtering strategy, since there is no overlap for different region masks within the same image.

3.3 Region-level Consistency Learning

Different from previous works [Zhong et al., 2021; Alonso et al., 2021] which used pixel-level consistency regularization, we develop a region-level consistency learning, which consists of a Region Class Consistency (RCC) Loss and Semantic Mask Consistency (SMC) loss. The former enforces the consistency of region classes, and the latter further promotes the consistency of the union regions with the same class.

Given the student class logits {psi ∈ △K+1 }Ni=1 and the pseudo segment label {cti ,mti }N t i=1 , the proposed Region Class Consistency loss is employed to enforce the class consistency of matched region pairs, which is defined as:

Furthermore, different regions may correspond to the same class, so we design a Semantic Mask Consistency loss to promote the consistency of the union regions with the same class and the pseudo semantic mask, which can be formulated as follow:

where Union(, ) means merging regions with the same class into a single region. Following DETR [Carion et al., 2020] and MaskFormer [Cheng et al., 2021], Lmask is a linear combination of a focal loss [Lin et al., 2017] and a dice loss [Milletari et al., 2016].

4 Experiment

4.1 Datasets

PASCAL VOC 2012. PASCAL VOC 2012 [Everingham et al., 2015] is a standard object-centric semantic segmentation dataset, which contains more than 13,000 images with 21 classes (20 object classes and 1 background class). In the original PASCAL VOC 2012 dataset (VOC Train), 1464 images are used for training, 1449 images for validation and 1456 images for testing. Following the common practice, we also use the augmented dataset (VOC Aug) which contains 10,582 images as the training set. For both the original and augmented datasets, 1/2, 1/4, 1/8, and 1/16 training images are used as labeled data, respectively, for conducting the semi-supervised experiments.

Cityscapes. Cityscapes [Cordts et al., 2016] dataset is designed for urban scene understanding with a high resolution (2048×1024). It contains 19 classes for scene semantic segmentation. The finely annotated 5,000 images that follow the official split have 2,975 images for training, 500 images for evaluation, and 1,525 images for testing. We employ 1/4 and 1/8 training images as the labeled data, and the remaining images are used as the unlabeled data.

Evaluationmetrics. We use single scale testing, and report the evaluation results of the model on Pascal VOC 2012 val set and Cityscapes val set by using mean of Intersection over Union (mIoU). We compare with state-of-the-art methods on different dataset partition protocols.

Implementation details. We use ResNet101 as our backbone, and use MaskFormer Head [Cheng et al., 2021] as segmentation head. We follow the original MaskFormer to set the hyper-parameters. For all experiments, we set the batch size to 16, and use ADAMW as the optimizer with an initial learning rate of 0.0001, and weight decay of 0.0001. Empirically, we set the loss weight of β1, β2, β3 and β4 to 1, 20, 4 and 4, respectively. In addition, we employ a poly learning rate policy which is multiplied by (1 - iter maxiter )power with power = 0.9. For PASCAL VOC 2012 dataset, we use the crop size of 512 × 512, and train our model for 120K and 160K iterations for VOC Train and VOC Aug, respectively. For Cityscapes dataset, we use the crop size of 768×768, and set training iterations as 120K without using any extra training data.

4.2 Comparison to the State-of-the-Arts

Results on PASCALVOC 2012. We show the comparison results on the PASCAL VOC 2012 val set in Table 1. All the models are trained using 4 V100 GPUs. On VOC Train, our proposed RC2L outperforms all existing pixel-level regularization methods, achieving the improvement of 1.18%, 0.53%, 1.45%, and 1.26% with partition protocols of 1/2, 1/4, 1/8, and 1/16 respectively. We also compare the results of the 1.4k/9k split, our RC2L is +5.18% higher than the most recent PC2Seg (79.33 vs. 74.15). On VOC Aug, RC2L also outperforms the previous state-of-the-art methods, obtaining the improvements of 1.79%, 2.03%, 1.05%, and 1.08% under 1/2, 1/4, 1/8, and 1/16 partitions, respectively.

Results on Cityscapes. Table 2 shows comparison results on Cityscapes val set. All the models are trained using 8 V100 GPUs. The performance of our method is improved by 2.53% and 2.51%, compared to the supervised baseline under partition protocols of 1/4 and 1/8, respectively. Our method also outperforms the previous state-of-the-art. For example, RC2L outperforms PC2Seg, which is the best pixellevel contrastive learning method by 1.32% and 1.75% under 1/4 and 1/8 partition protocols, respectively.

4.3 Ablation Study

Investigating each component. We investigate the effect of each component in our methods. The experimental results are illustrated in Table 3. The baseline model only uses mask consistency between teacher and student predictions for unlabeled data. It can be seen that the baseline method achieves 69.37% and 75.26% on 1/2 VOC Train and 1/4 VOC Aug datasets, respectively. We can see that SMC loss improves the baseline by 3.89% and 1.32% on 1/2 VOC Train and 1/4 VOC Aug datasets, respectively. RCC loss obtains the improvement of 2.81% and 1.31% over "Baseline + SMC loss". These experimental results demonstrate the effectiveness of our region-level consistency learning. Furthermore, RMC loss obtains the improvement of 0.78% and 1.03%. RFC loss further achieves the improvement of 0.21% and 0.79%. These experimental results demonstrate the effectiveness of our region-level contrastive learning.

Loss weight. We show the effect of loss weight α which is used to balance the supervised loss and unsupervised loss in Table 4. We find that α = 2 achieves the best performance on 1/2 VOC Train set. For 1/4 VOC Aug set, α = 1 obtains the best performance.

Number of queries. We study the number of queries on PASCAL VOC 2012 in Table 5. We can see that N = 50 achieves the best result on both 1/2 VOC Train and 1/4 VOC Aug, and N = 100 degrades the performance more than N = 20. This suggests that a small number of queries is sufficient to provide a good result for datasets with a few number of categories in just one image.

4.4 Visualization

Figure 3 shows the visualization results of our methods on PASCAL VOC 2012. We compare the predictions of our RC2L with the ground-truth, supervised baseline, and semisupervised consistency baseline. One can see that our RC2L can correct more noisy predictions compared to the supervised baseline and the semi-supervised consistency baseline. We follow CCT and directly complete consistency learning between student outputs and pseudo labels as the semisupervised consistency baseline. In particular, the supervised baseline mislabels some pixels in the 1st row and the 3rd row. Both the supervised baseline and the semi-supervised consistency baseline mistakenly classify some pixels in the 3rd row and the 5th row.

5 Conclusion

We have developed the Region-level Contrastive and Consistency Learning (RC2L) method for semi-supervised semantic segmentation. The core contributions of our RC2L are the proposed region-level contrastive and consistency regularization. The former consists of Region Mask Contrastive (RMC) loss and Region Feature Contrastive (RFC) loss, the latter contains Region Class Consistency (RCC) loss and Semantic Mask Consistency (SMC) loss. Extensive experiments on PASCAL VOC 2012 and Cityscapes have shown that our RC2L outperforms the state-of-the-art semi-supervised semantic segmentation methods, demonstrating that our regionlevel Contrastive and Consistency regularization can achieve better results than previous pixel-level regularization.