End-to-End Semi-Supervised Object Detection with Soft Teacher

Introduction

Data matters. In fact, large data such as ImageNet has largely triggered the boom of deep learning in computer vision. However, obtaining labels can be a bottleneck, due to the time-consuming and expensive annotation process. This has encouraged learning methods to leverage unlabeled data in training deep neural models, such as self-supervised learning and semi-supervised learning. This paper studies the problem of semi-supervised learning, in particular for object detection.

For semi-supervised object detection, we are concerned with the pseudo-label based approaches, which are the current state-of-the-art. These approaches conduct a multi-stage training schema, with the first stage training an initial detector using labeled data, followed by a pseudo-labeling process for unlabeled data and a re-training step based on the pseudo labeled unannotated data. These multi-stage approaches achieve reasonably good accuracy, however, the final performance is limited by the quality of pseudo labels generated by an initial and probably inaccurate detector trained using a small amount of labeled data.

To address this issue, we present an end-to-end pseudo-label based semi-supervised object detection framework, which simultaneously performs pseudo-labeling for unlabeled images and trains a detector using these pseudo labels along with a few labeled ones at each iteration. Specifically, labeled and unlabeled images are randomly sampled with a preset ratio to form one data batch. Two models are applied on these images, with one conducting detection training and the other in charge of annotating pseudo labels for unlabeled images. The former is also referred to as a student, and the latter is a teacher, which is an exponential moving average (EMA) of the student model. This end-to-end approach avoids the complicated multi-stage training scheme. Moreover, it also enables a “flywheel effect” that the pseudo labeling and the detection training processes can mutually reinforce each other, so that both get better and better as the training goes on.

Another important benefit of this end-to-end framework is that it allows for greater leverage of the teacher model to guide the training of the student model, rather than just providing “some generated pseudo boxes with hard category labels” as in previous approaches .

A soft teacher approach is proposed to implement this insight. In this approach, the teacher model is used to directly assess all the box candidates that are generated by the student model, rather than providing “pseudo boxes” to assign category labels and regression vectors to these student-generated box candidates. The direct assessment on these box candidates enables more extensive supervision information to be used in the student model training. Specifically, we first categorize the box candidates as foreground/background by their detection scores with a high foreground threshold to ensure a high precision of the positive pseudo labels, as in . This high foreground threshold, however, results in many positive box candidates mistakenly assigned as background. To address this issue, we propose using a reliability measure to weight the loss of each “background” box candidate. We empirically find that a simple detection score produced by the teacher model can well serve as the reliability measure, and is used in our approach. We find that this approach measure performs significantly better than previous hard foreground/background assignment methods (see Table 3 and Table 4), and we name it “soft teacher”.

Another approach instantiates this insight is to select reliable bounding boxes for the training of the student’s localization branch, by a box jittering approach. This approach first jitters a pseudo-foreground box candidate several times. Then these jittered boxes are regressed according the teacher model’s location branch, and the variance of these regressed boxes is used as a reliability measure. The box candidate with adequately high reliability will be used for the training of the student’s localization branch.

On MS-COCO object detection benchmark , our approach achieves 20.5 mAP, 30.7 mAP and 34.0 mAP on val2017 with 1%, 5% and 10% labeled data using the Faster R-CNN framework with ResNet-50 and FPN , surpassing previous best method STAC by +6.5, +6.4 and +5.4 mAP, respectively.

In addition, we also perform evaluation on a more challenge setting where the labelled data has been adequately large to train a reasonably accurate object detector. Specifically, we adopt the complete COCO train2017 set as labeled data and the unlabeled2017 set as the unlabeled data. Under this setting, we improve the supervised baseline of a Faster R-CNN approach with ResNet-50 and ResNet-101 backbones by +3.6 mAP and +3.0 mAP, respectively.

Moreover, on a state-of-the-art Swin-Transformer based detector which achieves 58.9 mAP for object detection and 51.2 mAP for instance segmentation on COCO test-dev2017, the proposed approach can still improve the accuracy by +1.5 mAP and +1.2 mAP, respectively, reaching 60.4 mAP and 52.4 mAP. Further incorporating with the Object365 pre-trained model, the detection accuracy reaches 61.3 mAP and the instance segmentation accuracy reaches 53.0 mAP, which is the new state-of-the-art on this benchmark.

Related works

Semi-supervised learning in image classification can be roughly categorized into two groups: consistency based and pseudo-label based. The consistency based methods leverage the unlabeled images to construct a regularization loss which encourages different perturbations of a same image to produce similar predictions. There are several ways to implement perturbations, including perturbing the model , augmenting the images or adversarial training . In , the training target is assembled by predicting different training steps. In , they develop by ensembling the model itself instead of the model prediction, the so-called exponential mean average (EMA) of the student model. The pseudo-label approaches (also named as self-training) annotate unlabeled images with pseudo labels by an initially trained classification model, and the detector is refined by these pseudo labeled images. Unlike our method which focuses on object detection, the pseudo-label does not have to solve the problem of assigning foreground/background labels and box regression when classifying images. Recently, some works explore the importance of data augmentation in semi-supervised learning, which inspire us to use the weak augmentation to generate pseudo-labels and the strong augmentation for the learning of detection models.

Semi-supervised learning in object detection

Similar to the semi-supervised learning in image classification, semi-supervised object detection methods also have two categories: the consistency methods and pseudo-label methods . Our method belongs to the pseudo-label category. In , the predictions of different data augmentation are ensembled to form the pseudo labels for unlabeled images. In , a SelectiveNet is trained to select the pseudo-label. In , a box detected on an unlabeled image is pasted onto a labeled image, and the localization consistency estimation is performed onto the pasted label image. As the image itself is modified, a very thorough detection process is required in . In our method, only the lightweight detection head is processed. STAC proposes to use a weak data augmentation for model training and a strong data augmentation is used for performing pseudo-label. However, like other pseudo-label methods , it also follows the multi-stage training scheme. In contrast, our method is an end-to-end pseudo-labeling framework, which avoids the complicated training process and also achieves better performance.

Object Detection

Object detection focuses on designing efficient and accurate detection framework. There are two mainstreams: single-stage object detectors and two-stage object detectors . The main difference between the two types of methods is whether to use a cascade to filter a large number of object candidates (proposals). In theory, our method is compatible with both types of methods. However, to allow a fair comparison with previous works on semi-supervised object detection, we use Faster R-CNN as our default detection framework to illustrate our method.

Methodology

Figure. 2 illustrates an overview of our end-to-end training framework. There are two models, a student model and a teacher model. The student model is learned by both the detection losses on the labeled images and on the unlabeled images using pseudo boxes. The unlabeled images have two sets of pseudo boxes, which are used to drive the training of the classification branch and the regression branch, respectively. The teacher model is an exponential moving average (EMA) of the student model. Within this end-to-end framework, there are two crucial designs: soft teacher and box jittering.

We first introduce the end-to-end framework for pseudo-label based semi-supervised object detection. Our approach follows the teacher-student training scheme. In each training iteration, labeled images and unlabeled images are randomly sampled according to a data sampling ratio $s_{r}$ to form a training data batch. The teacher model is performed to generate the pseudo boxes on unlabeled images and the student model is trained on both labeled images with the ground-truth and unlabeled images with the pseudo boxes as the ground-truth. Thus, the overall loss is defined as the weighted sum of supervised loss and unsupervised loss:

where $\mathcal{L}_{s}$ and $\mathcal{L}_{u}$ denote supervised loss of labeled images and unsupervised loss of unlabeled images respectively, $\alpha$ controls contribution of unsupervised loss. Both of them are normalized by the respective number of images in the training data batch:

where $I_{l}^{i}$ indicates the $i$-th labeled image, $I_{u}^{i}$ indicates the $i$-th unlabeled image, $\mathcal{L}_{\text{cls}}$ is the classification loss, $\mathcal{L}_{\text{reg}}$ is the box regression loss, $N_{l}$ and $N_{u}$ denote the number of labeled images and unlabeled images, respectively.

At the beginning of training, both the teacher model and student model are randomly initialized. As the training progresses, the teacher model is continuously updated by the student model, and we follow the common practices that the teacher model is updated by exponential moving average (EMA) strategy.

In contrast to taking a simple probability distribution as the pseudo-label in image classification, creating pseudo-label for object detection is more complicated since an image usually contains multiple objects and the annotation of objects consists of location and category. Given an unlabeled image, the teacher model is used to detect objects and thousands of box candidates are predicted. The non-maximum suppression (NMS) is then performed to eliminate redundancy. Although most redundant boxes are removed, there are still some non-foreground candidates left. Therefore, only candidates with the foreground scoreThe foreground score is defined as the maximum probability of all non-background categories. higher than a threshold are retained as the pseudo boxes.

In order to generate high-quality pseudo boxes and to facilitate the training of the student model, we draw on FixMatch which is the latest advancement in semi-supervised image classification task. Strong augmentation is applied for detection training of the student model and weak augmentation is used for pseudo-labeling of the teacher model.

In theory, our framework is applicable to mainstream object detectors, including single-stage object detectors and two-stage object detectors . To allow a fair comparison with previous methods, we use Faster R-CNN as our default detection framework to illustrate our method.

2 Soft Teacher

The performance of the detector depends on the quality of the pseudo-label. In practice, we find that using a higher threshold on foreground score to filter out most of the student-generated box candidates with low-confidence can achieve better results than using a lower threshold. As shown in Table. 9, the best performance is achieved when the threshold is set to 0.9. However, while the strict criteria (higher threshold) leads to higher foreground precision, the recall of the retained box candidates also falls off quickly. As shown in Figure. 3 (a), when the foreground threshold is set to 0.9, the recall is low, as $33\%$, while the precision reaches $89\%$. In this case, if we use IoU between student generated box candidates and teacher-generated pseudo boxes to assign foreground and background labels, as a general object detection framework does when real box annotations are provided, some foreground box candidates will be mistakenly assigned as negatives, which may hinder the training and harm the performance.

To alleviate this issue, we propose a soft teacher approach which leverages richer information from the teacher model, thanks to the flexibility of the end-to-end framework. Specifically, we assess the reliability of each student-generated box candidate to be a real background, which is then used to weigh its background classification loss. Given two box sets $\{b_{i}^{\text{fg}}\}$ and $\{b_{i}^{\text{bg}}\}$, with $\{b_{i}^{\text{fg}}\}$ denoting boxes assigned as foreground and $\{b_{i}^{\text{bg}}\}$ denoting the boxes assigned as background, the classification loss of an unlabeled image with the reliable weighting is defined as:

where $\mathcal{G_{\text{cls}}}$ denotes the set of (teacher-generated) pseudo boxes used for classification, $l_{\text{cls}}$ is the box classification loss, $r_{j}$ is the reliability score for j-th background box candidate, $N_{b}^{\text{fg}}$ and $N_{b}^{\text{bg}}$ are the number of box candidates of the box set $\{b_{i}^{\text{fg}}\}$ and $\{b_{i}^{\text{bg}}\}$, respectively.

Estimating the reliability score $r$ is challenging. We find empirically that the background score produced by the teacher model with weak augmented image can well serve as a proxy indicator of $r$ and is easily obtained in our end-to-end training framework. Specifically, given a student-generated box candidate, its background score can be obtained simply by using the teacher (BG-T) to process the box through its detection head. It is worth noting that this approach, unlike the widely used hard negative mining approaches, e.g., OHEM or Focal Loss , is more like a “simple” negative mining. For comparison, we also examine several other indicators:

Background score of student model (BG-S): Another natural way to generate the background score is to use the prediction of student model directly.

Prediction difference (Pred-Diff): The prediction difference between the student model and teacher model is also a possible indicator. In our approach, we simply use the difference between the background scores of the two models to define the reliability score:

where $p_{\text{S}}^{\text{bg}}$ and $p_{\text{T}}^{\text{bg}}$ are the predicted probability of the background class of the student and the teacher model, respectively.

Intersection-over-Union: The IoU between ground-truths and box candidates is a commonly used criterion for foreground/background assignment. There are two different hypotheses about how to use IoU to measure whether a box candidate belongs to the background. In the first hypothesis, if the IoU between a box candidate and a ground-truth box is less than a threshold (e.g., 0.5), a larger IoU indicates the box candidate has greater probability of being background. This can be viewed as an IoU-based hard negative mining which is adopted by Fast R-CNN and Faster R-CNN in the early implementation. In contrast, the other hypothesis suggests that box candidates with a smaller IoU with ground-truths are more likely to be backgrounds. In our experiments, we validate both hypotheses and name them as IoU and Reverse-IoU.

3 Box Jittering

As shown in Figure. 3 (b), the localization accuracy and the foreground score of the box candidates do not show a strong positive correlation, which means that the boxes with high foreground score may not provide accurate localization information. This indicates that the selection of the teacher-generated pseudo boxes according to the foreground score is not suitable for box regression, and a better criterion is needed.

We introduce an intuitive approach to estimate the localization reliability of a candidate pseudo box by measuring the consistency of its regression prediction. Specifically, given a teacher-generated pseudo box candidate $b_{i}$, we sample a jittered box around $b_{i}$ and feed the jittered box into the teacher model to obtain the refined box $\hat{b}_{i}$, which is formulated as follows:

The above procedure is repeated several times to collect a set of $N_{\text{jitter}}$ refined jittered boxes $\{\hat{b}_{i,j}\}$, and we define the localization reliability as the box regression variance:

where $\sigma_{k}$ is the standard derivation of the k-th coordinate of the refined jittered boxes set $\{\hat{b}_{i,j}\}$, $\hat{\sigma_{k}}$ is the normalized $\sigma_{k}$, $h(b_{i})$ and $w(b_{i})$ represent the height and width of box candidate $b_{i}$, respectively.

A smaller box regression variance indicates a higher localization reliability. However, computing the box regression variances of all pseudo box candidates is unbearable during training. Therefore, in practice, we only calculate the reliability for the boxes with a foreground score greater than 0.5. In this way, the number of boxes that need to be estimated is reduced from an average of hundreds to around 17 per image and thus the computation cost is almost negligible.

In Figure. 3 (c), we illustrate the correlation between the localization accuracy and our box regression variance. Compared with the foreground score, the box regression variance can better measure the localization accuracy. This motivates us to select box candidates whose box regression variance is smaller than a threshold as pseudo-label to train the box regression branch on unlabeled images. Given the pseudo boxes $\mathcal{G}_{\text{reg}}$ for training the box regression on unlabeled data, the regression loss is formulated as:

where $b_{i}^{\text{fg}}$ is i-th box assigned as foreground, $N_{b}^{\text{fg}}$ is the total number of foreground box, $l_{\text{reg}}$ is the box regression loss. Therefore, by substituting Equ. 4 and Equ. 10 into Equ. 3, the loss of unlabeled images is:

Here we use the pseudo boxes $\mathcal{G}_{\text{cls}}$ and $\mathcal{G}_{\text{reg}}$ as the inputs of the loss to highlight the fact that the pseudo boxes used in classification and box regression are different in our approach.

Experiments

We validate our method on the MS-COCO benchmark . Two training datasets are provided, the train2017 set contains 118k labeled images and the unlabeled2017 set contains 123k unlabeled images. In addition, the val2017 set with 5k images is also provided for validation. In previous methods , there are two settings for validating the performance:

Partially Labeled Data: STAC first introduced this setting. 1%, 5% and 10% images of train2017 set are sampled as the labeled training data, and the remaining unsampled images of train2017 are used as the unlabeled data. For each protocol, STAC provides 5 different data folds and the final performance is the average of all 5 folds.

Fully Labeled Data: In this setting, the entire train2017 is used as the labeled data and unlabeled2017 is used as the additional unlabeled data. This setting is more challenging. Its goal is to use the additional unlabeled data to improve a well-trained detector on large-scale labeled data.

We evaluate our method on both settings and follow the convention to report the performance on val2017 with the standard mean average precision (mAP) as the evaluation metrics.

2 Implementation Details

We use the Faster R-CNN equipped with FPN (Feature Pyramid Network) as our default detection framework to evaluate the effectiveness of our method, and an ImageNet pre-trained ResNet-50 is adopted as the backbone. Our implementation and hyper-parameters are based on MMDetection . Anchors with 5 scales and 3 aspect ratios are used. 2k and 1k region proposals are generated with a non-maximum suppression threshold of 0.7 for training and inference. In each training step, 512 proposals are sampled from 2k proposals as the box candidates to train RCNN. Since the amount of training data of Partially Labeled Data setting and Full Labeled Data setting has large differences, the training parameters under the two settings are slightly different.

Partially Labeled Data: The model is trained for 180k iterations on 8 GPUs with 5 image per GPU. With SGD training, the learning rate is initialized to 0.01 and is divided by 10 at 110k iteration and 160k iteration. The weight decay and the momentum are set to 0.0001 and 0.9, respectively. The foreground threshold is set to 0.9 and the data sampling ratio $s_{r}$ is set to 0.2 and gradually decreases to 0 over the last 10k iterations.

Fully Labeled Data: The model is trained for 720k iterations on 8 GPUs with 8 image per GPU. In SGD training, the learning rate is initialized to 0.01 and is divided by 10 at 480k iteration and 680k iteration. The weight decay and the momentum are set to 0.0001 and 0.9, respectively. The foreground threshold is set to 0.9 and the data sampling ratio $s_{r}$ is set to 0.5 and gradually decreases to 0 in the last 20k iterations.

For estimating the box localization reliability, we set $N_{\text{jitter}}$ as 10, and threshold is set as 0.02 to select the pseudo-labels for box regression. 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. In addition, we follow STAC and FixMatch to use different data augmentation for pseudo-label generation, labeled image training and unlabeled image training. The details are summarized in Table .2.

3 System Comparison

In this section, we compare our method with previous state-of-the-arts on MS-COCO. We first evaluate on the Partially Labeled Data setting and compare our method with STAC. For benchmarking, we compare the supervised baseline of our method with the results reported in STAC and find they perform similarly, the results are shown in Table. 3. In this case, we further compare our method with STAC at the system level, and our method shows a significant performance improvement in different protocols. Specifically, our method outperforms the STAC by 6.5 points, 6.4 points and 5.4 points when there are $1\%$, $5\%$, and $10\%$ labeled data, respectively. The qualitative results of our method compared with supervised baseline are shown in Figure. 4.

Then we compare our method with other state-of-the-art methods in Fully Labeled Data setting. Since the reported performance of supervised baseline varies in different works, we report the results of the comparison methods and their baseline at the same time. The results are shown in Table. 4.

We first compare with the Proposal Learning and STAC which also use unlabeled2017 as additional unlabeled data. Because of the better hyper-parameters and more adequate training, our supervised baseline achieved better performance than other methods. Under the stronger baseline, our method still shows a greater performance gain (+3.6 points) than Proposal Learning (+1.0 points) and STAC (-0.3 points). Self-training uses ImageNet (1.2M images) and OpenImages (1.7M images) as the additional unlabeled data, which is 20$\times$ larger than the unlabeled2017 (123k images) that we use. With similar baseline performance, our method also shows better result with less unlabeled data.

In addition, we further evaluate our method on other stronger detectors, and the results evaluated on val2017 set are shown in Table. 5. Our method consistently improves the performance of different detectors by a notable margin. Even in the state-of-the-art detector HTC++ with Swin-L backbone, we still show 1.8 improvement on detection AP and 1.4 improvement on mask AP. Moreover, we also report the results on test-dev2017 set. As shown in Tabel. 1, our method improves the HTC++ with Swin-L backbone by 1.5 mAP on detection, which is the first work to surpass 60 mAP on COCO object detection benchmark.

4 Ablation Studies

In this section, we validate our key designs. If not specified, all the ablation experiments are performed on the single data fold provided by with 10% labeled images from train2017 set.

We compare our end-to-end method with the multi-stage framework as shown in Table 6. By simply switching from the multi-stage framework to our end-to-end framework, performance is increased by 1.3 points. By updating the teacher model with the student model through the exponential moving average (EMA) strategy, our method further achieves 31.2 mAP.

Effects of Soft Teacher and Box Jittering.

We ablate the effects of soft teacher and box jittering. The results are shown in Table. 7. Based on our end-to-end model equipped with EMA (E2E+EMA), integrating the soft teacher improves the performance by 2.4 points. Further applying the box jittering, the performance reaches 34.2 mAP, which is 3 points better than E2E+EMA.

Different Indicators in Soft Teacher.

In Section. 3.2, several different indicators are explored for reliability estimation. Here, we evaluate the different indicators and the results are shown in Table. 8. The background score predicted by the teacher model achieves the best performance. Simply switching the model from teacher to student will make the performance worse. In addition, the improvement of IoU and Revearse-IoU is negligible compared with BG-T. These results prove the necessity of leveraging the teacher model.

Effects of other hyper-parameters.

We study the effects of hyper-parameters used in our method. Table. 9 studies the effects of different foreground score thresholds. The best performance is achieved when the threshold is set to 0.9, and lower or higher thresholds will cause significant performance degradation. In Table. 10, we study the box regression variance threshold. The best performance is achieved when the threshold is set to 0.02. In Table. 11, we study the effects of different number of jittered boxes, and the performance is saturated when $N_{\text{jitter}}$ is set to 10.

Conclusion

In this paper, we proposed an end-to-end training framework for semi-supervised object detection, which discards the complicated multi-stage schema adopted by previous approaches. Our method simultaneously improves the detector and pseudo labels by leveraging a student model for detection training, and a teacher model which is continuously updated by the student model through the exponential moving average strategy for online pseudo-labeling. Within the end-to-end training, we present two simple techniques named soft teacher and box jittering to facilitate the efficient leverage of the teacher model. The proposed framework outperforms the state-of-the-art methods by a large margin on MS-COCO benchmark in both partially labeled data and fully labeled data settings.

Acknowledgement

We would like to thank Yue Cao for his valuable suggestions and discussions; Yutong Lin and Yixuan Wei for help on Swin Transformer experiments.

References