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

The overall framework of our RC²L 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 (with ground-truth) and unlabeled data.

Specifically, given an image-label pair, $\{x_{l},z^{gt}\}$, $x_{l}$ is the image, and $z^{gt}$ is the corresponding annotations that are the set of $N^{gt}$ ground truth segments, i.e., $z^{gt}=\{(c_{i}^{gt},m_{i}^{gt})|c_{i}^{gt}\in\{1,2,...,K\},m_{i}^{gt}\in\{0,1\}^{H\times W}\}_{i=1}^{N^{gt}}$, where $c_{i}$ is the ground truth class of the $i^{th}$ segment, $W$ and $H$ represent the width and height of the input image. We feed the image $x_{l}$ into the student model, outputting probability-mask pairs $z^{l}=\{(p_{i}^{l},m_{i}^{l})|p_{i}^{l}\in\triangle^{K+1},m_{i}^{l}\in[0,1]^{H\times W}\}_{i=1}^{N}$, where the probability distribution $p_{i}$ contains a “no object” label $\oslash$ and $K$ category labels. A bipartite matching-based assignment $\sigma$ between the set of predictions $z^{l}$ and $z^{gt}$ is conducted for computing supervised loss:

$$\begin{split}\mathcal{L}^{label}(z^{l},z^{gt})=&\sum_{i=1}^{N}[-logp_{\sigma(i)}^{l}(c_{i}^{gt})+\\ &1_{c_{i}^{gt}\neq\oslash}\mathcal{L}_{mask}(m_{\sigma(i)}^{l},m_{i}^{gt})],\end{split}$$

(1)

where $\mathcal{L}_{mask}$ is a binary mask loss. Please refer to Cheng et al. (2021) for more details.

For the unlabeled image $x_{u}$, we feed its weak augmentation $x_{u}^{w}$ and strong augmentation $x_{u}^{s}$ 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 model is employed to generate pseudo label $z^{t}=\{(c_{i}^{t},m_{i}^{t})\}_{i=1}^{N^{t}}$ which is used to guide the training of the student model. The unsupervised loss can be formulated as follow:

$$\begin{split}\mathcal{L}^{unlabel}(x_{u})=\beta_{1}\mathcal{L}_{RCC}+\beta_{2}\mathcal{L}_{SMC}+\\ \beta_{3}\mathcal{L}_{RMC}+\beta_{4}\mathcal{L}_{RFC},\end{split}$$

(2)

where $\mathcal{L}_{RCC}$ and $\mathcal{L}_{SMC}$ mean the Region Class Consistency (RCC) loss and Semantic Mask Consistency (SMC) loss (Section 3.2), respectively. $\mathcal{L}_{RMC}$ and $\mathcal{L}_{RFC}$ denote the Region Mask Contrastive (RMC) loss and Region Feature Contrastive (RFC) loss (Section 3.3), respectively. $\beta_{1}$, $\beta_{2}$, $\beta_{3}$ and $\beta_{4}$ denote the loss weight. Therefore, the total loss is defined as:

$$\mathcal{L}=\mathcal{L}^{label}(z^{l},z^{gt})+\alpha\mathcal{L}^{unlabel}(x_{u}),$$

(3)

where $\alpha$ is a constant to balance between the supervised and unsupervised losses.

Following Mean Teacher Tarvainen and Valpola (2017), the parameters $\theta_{t}$ of teacher model are an exponential moving average of parameters $\theta_{s}$ of the student model. Specifically, at every training step, the parameters $\theta_{t}$ of teacher network is updated as follows:

$$\theta_{t}=\tau\theta_{t}+(1-\tau)\theta_{s},$$

(4)

where $\tau\in\ [0,1\ ]$ is a decay rate.

The goal of region-level contrastive learning is to increase the similarity between each region mask and feature from strong augmentation $x_{u}^{s}$ and the mask and feature of the matched region (from weak augmentation $x_{u}^{w}$), 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 $x_{u}^{w}$ is firstly fed into the teacher network, which outputs pseudo labels $z^{t}=\{(c_{i}^{t},m_{i}^{t})\}_{i=1}^{N^{t}}$. The strong augmentation $x_{u}^{s}$ is fed into the student network, which outputs class logits $p^{s}=\{p_{i}^{s}\in\triangle^{K+1}\}_{i=1}^{N}$, mask embeddings $f^{s}=\{f_{i}^{s}\in\mathcal{R}^{C}\}_{i=1}^{N}$ , and per-pixel features $F^{s}\in\mathcal{R}^{C\times H/4\times W/4}$. Firstly, we obtain binary mask predictions $\{m_{i}^{s}\}_{i=1}^{N}$ via a dot product between the mask embeddings and per-pixel features. Then, the bipartite matching-based assignment $\sigma$ between student predictions $\{m_{i}^{s}\}_{i=1}^{N}$ and pseudo segment set $\{m_{i}^{t}\}_{i=1}^{N^{t}}$ is used for getting matched index set $ID=\{\sigma(i)\}_{i=1}^{N^{t}}$ and computing RMC Loss:

$$\mathcal{L}_{RMC}=\sum_{i}^{N^{t}}-log\frac{exp^{d(m_{\sigma(i)}^{s},m_{i}^{t})/\tau_{m}}}{\sum_{\{j\in ID,j\neq\sigma(i)\}}^{N}exp^{d(m_{j}^{s},m_{i}^{t})/\tau_{m}}},$$

(5)

where $\tau_{m}$ is a temperature hyper-parameter to control the scale of terms inside exponential, and $d(m_{i}^{s},m_{j}^{t})=\frac{2|m_{i}^{s}\cap m_{j}^{t}|}{|m_{i}^{s}|\cup|m_{j}^{t}|}$ measures the similarity between two masks.

Next, we compute region features $r^{s}=\{r_{\sigma(i)}^{s}\}_{i=1}^{N^{t}}$ via combining per-pixel features $F^{s}$ and region mask set $\{m_{i}^{s}\}_{i=1}^{N}$. The process can be formulated as follow:

$$r_{\sigma(i)}^{s}=GAP(m_{\sigma(i)}^{s}\cdot F^{s}),\forall\sigma(i),$$

(6)

where “GAP” indicates a global average pooling operation. Similarly, we compute target region features $r^{t}=\{r_{i}^{t}\}_{i=1}^{N^{t}}$ as follow:

$$r_{i}^{t}=GAP(m_{i}^{t}\cdot F^{s}),\forall i,$$

(7)

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:

$$\mathcal{L}_{RFC}=\sum_{i}^{N^{t}}-log\frac{exp^{cos(r_{\sigma(i)}^{s},r_{i}^{t})/\tau_{f}}}{\sum_{\{j\in ID,j\neq\sigma(i)\}}^{N}exp^{cos(r_{j}^{s},r_{i}^{t})/\tau_{f}}},$$

(8)

where $\tau_{f}$ is a temperature hyper-parameter to control the scale of terms inside exponential, $cos(u,v)=\frac{u^{\mathrm{T}}v}{\lVert u\rVert\lVert v\rVert}$ is the cosine similarity.

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 $\{p_{i}^{s}\in\triangle^{K+1}\}_{i=1}^{N}$ and the pseudo segment label $\{c_{i}^{t},m_{i}^{t}\}_{i=1}^{N^{t}}$, the proposed Region Class Consistency loss is employed to enforce the class consistency of matched region pairs, which is defined as:

$$\mathcal{L}_{RCC}=\sum_{i}^{N^{t}}-log\,p_{\sigma(i)}^{s}(c_{i}^{t}),$$

(9)

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:

$$\mathcal{L}_{SMC}=\sum_{i}^{N^{t}}\mathcal{L}_{mask}(Union(\{m_{i}^{s}\}_{i=1}^{N},c_{i}^{t}),m_{i}^{t}),$$

(10)

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), $\mathcal{L}_{mask}$ is a linear combination of a focal loss Lin et al. (2017) and a dice loss Milletari et al. (2016).

Figure 3 shows the visualization results of our methods on PASCAL VOC 2012. We compare the predictions of our RC²L with the ground-truth, supervised baseline, and semi-supervised consistency baseline. One can see that our RC²L 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 semi-supervised 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.