Revisiting Class Imbalance for End-to-end Semi-Supervised Object Detection

This paper aims to perform robust pseudo-label-based end-to-end SSOD where a set of labeled images $D_{l}={\{x_{l,i};y_{l,i}\}}^{N_{l}}_{i=1}$ and a set of unlabeled images $D_{u}={\{x_{u,j}\}}^{N_{u}}_{j=1}$ are used for training. Here, $N_{l}$ and $N_{u}$ are the number of labeled and unlabeled images. Further, $x_{l}$ and $y_{l}$ denote the image and its ground-truth annotations i.e., class labels and bounding box coordinates, respectively.

The architectural pipeline of proposed end-to-end network is illustrated in Figure 3. In each training iteration, labeled and unlabeled images are arbitrarily sampled to form an input batch. The teacher network produces the pseudo labels based on the weak and strong augmented unlabeled images. While the student network is trained using weakly augmented labeled images with ground-truth and strongly augmented unlabeled images with pseudo-labels as ground-truth. Here, the student network is trained using the weighted combination of the supervised and pseudo-label loss. This can be expressed mathematically as

$$L_{Total}=L_{sup}+\lambda\cdot(L^{wa}_{pl}+L^{sa}_{pl}).$$

(1)

Here, $\lambda$ controls the contribution of pseudo label loss. $L_{sup}$ is the supervised loss consisting $L^{t}_{sup}$ and $L^{s}_{sup}$ loss associated with the teacher and student network, respectively. While $L^{wa}_{pl}$ and $L^{sa}_{pl}$ are the pseudo-label losses based on the weak and strong augmented samples, respectively. These losses are mathematically described as

$$L_{sup}=\frac{1}{N_{l}}\sum_{i=1}^{N_{l}}\Big{(}L^{cls}_{sup}(I_{l,i})+L^{reg}_{sup}(I_{l,i})\Big{)},$$

(2)

$$L^{wa}_{pl}=\frac{1}{N_{u}}\sum_{i=1}^{N_{u}}\Big{(}L^{cls}_{pl}(I_{u,i}^{wa})+L^{reg}_{pl}(I_{u,i}^{wa})\Big{)},$$

(3)

$$L^{sa}_{pl}=\frac{1}{N_{u}}\sum_{i=1}^{N_{u}}\Big{(}L^{cls}_{pl}(I_{u,i}^{sa})+L^{reg}_{pl}(I_{u,i}^{sa})\Big{)}.$$

(4)

Here, $I_{l,i}$ denotes the $i^{th}$ labeled image. $I_{u,i}^{sa}$ and $I_{u,i}^{wa}$ indicates the $i^{th}$ strong augmented and weak augmented unlabeled image, respectively. $L^{cls}_{sup}$ and $L^{reg}_{sup}$ are the supervised classification and regression loss. Similarly, $L^{cls}_{pl}$ and $L^{reg}_{pl}$ are pseudo-label based classification and regression losses. Here, the number of labeled images and unlabeled images are noted as $N_{l}$ and $N_{u}$, respectively.

During the training process, the student network is trained using the weighted loss function given in Eq.1, and the teacher network is updated via the Double Exponential Moving Average (DEMA) update mechanism [19]. The teacher network predicts many bounding boxes for an unlabeled image. Hence, we employ the Non-Max Suppression (NMS) to eliminate redundancy. Although most of the redundant boxes are removed, some non-foreground candidates may remain. Therefore, only candidates with a foreground score¹¹1The foreground score is defined as the maximum probability of all non-background categories. greater than an adaptive threshold are retained as pseudo boxes. These pseudo boxes are then utilized in the classification loss. To learn box regression, the bounding boxes are passed through our Jitter-Bagging module to select reliable pseudo boxes, which are subsequently refined by the adaptive threshold filter.

In the following subsections, we discuss the adaptive threshold filter, efficient classification loss, Jitter-Bagging module and update mechanism in detail.

The performance of the detector network depends on quality of the pseudo-label. However, quality degrades due to the class imbalance, especially when there are few annotations. For underrepresented classes, the teacher network produces relatively lower confidence score [3], which barely survives the large threshold $\tau$. On the other hand, simply lowering $\tau$ leads to noisier pseudo-labels in common classes. Therefore, we propose adaptive threshold filter that adjusts the threshold value based on the confidence scores of the background and foreground bounding boxes for each category. The adaptive threshold (i.e., $\tau_{a}$) is mathematically defined for $N_{b}^{fg}$ number of foreground and $N_{b}^{bg}$ number of background bounding boxes as

$$\tau_{a}=\Bigg{\lfloor}\Bigg{(}\frac{\frac{1}{N_{b}^{fg}}\sum_{i=1}^{N_{b}^{fg}}\mathcal{S}_{i}^{fg}}{\frac{1}{N_{b}^{bg}}\sum_{j=1}^{N_{b}^{bg}}\mathcal{S}_{j}^{bg}}\Bigg{)}^{\gamma}\Bigg{\rfloor},$$

(5)

where $\gamma$ controls the degree of the underrepresented classes and it is set to 0.05, $\lfloor\cdot\rfloor$ indicates the closest decimal floor function for single precision (e.g. 0.94 will be set to 0.9). Here, $\mathcal{S}_{i}^{fg}$ and $\mathcal{S}_{j}^{bg}$ denote the scores obtained from the $i^{th}$ and $j^{th}$ foreground and background bounding boxes.

At the beginning of the training, when parameters of the networks are not fully learned, the majority is mispredictions, i.e., predominantly background predictions than foreground predictions. So adaptive threshold outputs a relatively smaller value and as the training progresses towards convergence, foreground predictions increases and thus, the adaptive threshold provides a larger value and more constraints for selecting appropriate bounding boxes.

For classification task, the overall loss is obtained by combining four different losses: foreground classification loss (i.e., $L^{cls}_{fg}$), background classification loss (i.e., $L^{cls}_{bg}$), background similarity loss (i.e., $L_{bg}^{sim}$) and foreground-background dissimilarity loss (i.e., $L_{fg-bg}^{dissim}$). The overall classification loss function can be defined as

$$L_{pl}^{cls}=L^{cls}_{fg}+L^{cls}_{bg}+L_{bg}^{sim}+L_{fg-bg}^{dissim}.$$

(6)

1) Foreground classification loss: Given student-generated foreground bounding boxes (i.e., $b^{fg}$), the foreground classification loss is defined as

$$L^{cls}_{fg}=\frac{1}{N_{b}^{fg}}\sum_{i=1}^{N_{b}^{fg}}l_{cls}(b_{i}^{fg}(s),\mathcal{B}_{cls}),$$

(7)

where $\mathcal{B}_{cls}$ denotes the set of teacher-generated pseudo boxes used for classification, $l_{cls}$ is the box classification loss²²2We use standard cross-entropy loss function as classification loss., $N_{b}^{fg}$ is the number of box candidates of box set ${b^{fg}}$.
2) Background classification loss: We employ the same loss function proposed by Xu et al. [32] for the background classification loss. Given background bounding boxes (i.e., ${b^{bg}}$), the classification loss is calculated as

$$L^{cls}_{bg}=\sum_{j=1}^{N_{b}^{bg}}\delta_{j}l_{cls}(b_{j}^{bg}(s),\mathcal{B}_{cls})$$

(8)

where, $\delta_{j}$ denoted as reliability weighting factor associated with $j^{th}$ sample and the same can be expressed as

$$\delta_{j}=\frac{r_{j}}{\sum_{k=1}^{N_{b}^{bg}}{r_{k}}}$$

(9)

Here $r_{j}$ is the reliability score for $j^{th}$ background box candidate, $N_{b}^{bg}$ is the number of box candidates of box set ${b^{bg}}$.
3) Background similarity loss: Inspired from [18], we introduce a novel loss to match the scores obtained from the background bounding boxes generated through teacher and student networks. Minimizing this loss will ensure that the teacher-generated pseudo boxes and student-predicted boxes become as similar as possible. The background similarity loss can be expressed as follows:

$$L_{bg}^{sim}=\frac{1}{N_{b}^{bg}}\sum_{i=1}^{N_{b}^{bg}}\beta\cdot\log(|e^{|\mathcal{S}_{i}^{bg}(s)|}-e^{|\mathcal{S}_{i}^{bg}(t)|}|+1),$$

(10)

where, $\beta$ denotes the controlling parameter, $\mathcal{S}_{i}^{bg}(t)$ and $\mathcal{S}_{i}^{bg}(s)$ indicates $i^{th}$ scores obtained from the background bounding boxes generated using teacher and student networks, respectively.
4) Foreground-Background dissimilarity loss: Inspired from relativistic average discriminator [7], a novel loss is introduced to separate out the foreground and background bounding boxes. The introduced loss considers the dissimilarity between scores obtained from background and foreground bounding boxes. Mathematically, this can be addressed as follows:

$$L_{fg-bg}^{dissim}=\frac{1}{N_{b}^{fg}}\sum_{i=1}^{N_{b}^{fg}}\big{(}1-|\mathcal{S}_{i}^{fg}(s)-\frac{1}{N_{b}^{bg}}\sum_{j=1}^{N_{b}^{bg}}\mathcal{S}_{j}^{bg}(s)|\big{)}.$$

(11)

Here, $\mathcal{S}_{i}^{fg}(s)$ and $S_{i}^{bg}(s)$ indicate $i^{th}$ and $j^{th}$ scores obtained from the student-generated foreground and background bounding boxes.

Xu et al. [32] found that the selection of the teacher-generated pseudo boxes according to the foreground score is not suitable for box regression. To tackle this issue, we introduce a Jitter-Bagging module where we sample a jittered-box around the teacher-generated pseudo box candidate $b_{i}$ and then feed to the teacher network to obtain refined box $b_{i}^{{}^{\prime}}$. This can be formulated as

$$b_{i}^{{}^{\prime}}=f_{Jitter}(b_{i}).$$

(12)

Here, $f_{Jitter}$ denotes the function of the Jitter operation. This procedure is repeated several times to collect the set of $N_{jitter}$ refined jittered boxes (i.e., $\{b_{i}^{{}^{\prime}}\}$). These refined jittered boxes are then fed to the traditional bagging algorithm which helps to obtain optimum refined boxes. Mathematically, this can be stated as:

$$\hat{b}_{i}=f_{Bagging}(\{b_{i}^{{}^{\prime}}\}),$$

(13)

where, $f_{Bagging}=max(\cdot)$ is the bagging operation which selects the bounding box candidate with maximum area. The obtained bounding boxes (i.e., $\hat{b}_{i}$) are further passed through an adaptive threshold filter to generate foreground bounding boxes (i.e., $\hat{b}_{i}^{fg}$). Finally, given the pseudo boxes $\mathcal{B}_{reg}$ for training the box regression on unlabeled data, the regression loss is formulated as

$$L_{pl}^{reg}=\frac{1}{N_{b}^{fg}}\sum_{i=1}^{N_{b}^{fg}}l_{reg}(\hat{b}_{i}^{fg},\mathcal{B}_{reg}),$$

(14)

where, $N_{b}^{fg}$ is the total number of foreground box, $l_{reg}$ is the box regression loss³³3We use standard mean absolute error for regression task..

At each iteration, the teacher weights get marginal updates using the student’s weights via the update mechanism. The gradually updated teacher is prone to weight fluctuations of the student network when a teacher network mispredicts a label. In [32, 25], authors used the Exponential Moving Average (EMA) mechanism to mitigate the negative effect of incorrect pseudo-labels [26]. However, the EMA update mechanism faces an issue of lagging, which limits its performance during sudden weight fluctuations. In this work, we employ EMA’s extension, i.e., Double Exponential Moving Average (DEMA) [19, 27]. The DEMA provides higher weight to the most recent observations and removes the inherent lag as compared to the EMA update mechanism. Further, the DEMA update mechanism can be mathematically defined as:

$$w(t)_{ts}=2\cdot EMA(t)_{ts}-EMA(EMA(t))_{ts},$$

(15)

where, the $EMA$ can be expressed as

$$EMA(t)_{ts}=\alpha^{2}w(t)_{ts-1}+(1-\alpha^{2})w(s)_{ts}.$$

(16)

Here, $w(t)_{ts}$ and $w(s)_{ts}$ are the weights of teacher and student network at current timestamp $ts$ and $\alpha$ is set to $0.999$.

We present our results on benchmark MS-COCO [13] and Pascal VOC [4] datasets.

Furthermore, we follow the data augmentation guidelines given in STAC [23] and FixMatch [22] for training and pseudo-label generation. For evaluation, we use the mean average precision (mAP)⁵⁵5Average Precision (AP) is finding the area under precision-recall curve. The mAP is defined as the average of AP. metric with its different variants i.e., mAP at IoU=0.5 (mAP@50), IoU=0.75 (mAP@75) and IoU=0.5:0.95 (mAP).

All experiments are performed with NVIDIA A100 dual GPUs. To allow a fair comparison with previous methods, we use Faster R-CNN [21] as our default detection network with pre-trained ResNet-50 [5] as a backbone. For training and inference, 2k and 1k region proposals have been generated using a non-maximum suppression threshold of 0.7. We sample 512 out of 2k proposals as the box candidates in each training step.

The proposed network is trained for 180k with 0.2 data sampling ratio and 720k iterations with 0.5 data sampling ratio for partially labeled data setting and fully labeled data setting, respectively. We adopt SGD as optimizer with learning rate of 0.001 divided by 10 at 120k and 160k iterations for partially labeled data setting and at 480k and 680k iterations for fully labeled data setting. Initially, we set foreground threshold to 0.5 and then it adjusts itself adaptively between 0.5 to 0.9.

We set $N_{jitter}$ to 10 in the proposed Jitter-Bagging module for regression task. The jittered boxes are randomly sampled by adding the offsets on four coordinates, and the offsets are uniformly sampled from [-6%, +6%] of the height or width of the pseudo box candidates.

This section provides the result comparison between the proposed network along with its supervised network on MS-COCO dataset. The validation performance is tabulated in Table 1. Here, one can see that our proposed network shows a significant performance improvement than the supervised baseline network in all protocols. In addition, we present the visual comparison for different proportion of labeled data in Figure 4. Here, it is clearly observed that the proposed network is able to detect tiny and occluded objects with better confidence score than that of the supervised framework.

All experiments are carried out on 10% partially labeled setting of MS-COCO. However, the analysis on 1% and 5% data settings are covered in the supplementary material.

Effectiveness of adaptive threshold filter: To prove the effectiveness of proposed adaptive threshold filter, few experiments have been carried out and the corresponding results are presented in Table 5. Here, we can see that the proposed adaptive threshold filter performs better than the static threshold values. To check its effectiveness over the threshold module proposed by Li et al. [11], we employed their thresholding module in our framework. The corresponding results are added in Table 5, which is marginally inferior to the proposed thresholding module. In our adaptive mechanism, we have used discrete thresholding to reduce fluctuations in the threshold value, Further, we trained a variant of our model with a continuous form of threshold and found lower performance than proposed discrete form.

Importance of Jitter-Bagging module: The ablation analysis of Jitter-Bagging module is presented in Table 6, where we can see that the proposed Jitter-Bagging module achieves highest performance and shows +1.2% absolute improvement in mAP measure over without Jitter-Bagging module. Interestingly, when we employ the Box Jittering [32] in our network, we observe that the proposed Jitter-Bagging still obtains +0.6% higher mAP than Box Jittering.

Effect of losses for classification: For classification task, we introduce two novel losses; background similarity loss i.e., $L_{bg}^{sim}$ and foreground-background dissimilarity loss i.e., $L_{fg-bg}^{dissim}$. To check its effectiveness, the proposed network is trained without both introduced losses (i.e., Case I), with background similarity loss (i.e., Case II) and with foreground-background dissimilarity loss (i.e., Case III). The corresponding measures are depicted in Table 7. Here, it can be noticed that both introduced losses help the proposed network to obtain better mAP measures. Additionally, Figure 5 shows the effect of loss values during the training iteration. Here, it can be seen that the proposed network with both losses converges better than others.

Effect of update mechanism: To verify the effectiveness of the DEMA update mechanism, we ablate the proposed network trained using EMA mechanism as well as employing deep copy configuration (i.e., weights of teacher network are copied from student network). The corresponding results are noted in Table 8, where it can be seen that DEMA helps to obtain +1.2% higher mAP measure, demonstrating its efficacy over EMA update mechanism.

Importance of label generator module: In our proposed network, we use two label generator modules; one associated with weak augmented samples while the other is based on strong augmented sample. To see the importance of this setting, the proposed network with individual label generator is trained and the obtained results are presented in Table 9. Here, it is observed that the proposed network with both label generators outperforms the individual label generator settings.