Unbiased Teacher for Semi-Supervised Object Detection

In plain words

Labeling objects with boxes is slow, so an object detector usually trains on a small labeled set. This method also lets the detector learn from unlabeled photos. A "teacher" copy of the detector looks at an unlabeled image and guesses which boxes contain objects. A "student" learns from those guesses, but it sees a strongly altered version of the same image. The teacher is never trained directly. It is a slowly moving average of the student, so its guesses change gradually from step to step.

Detection has a complication: background dominates the image regions, and some object classes are far more common than others. A teacher trained on little data tends to guess the common classes and the background. The method reweights the classification loss so that hard, uncertain examples count more, and it relies on the slow teacher update to keep the guesses stable.

The authors report large gains over earlier methods when only 0.5% to 2% of the COCO boxes are labeled. These are their own numbers.

TL;DR

• Semi-supervised object detection for Faster R-CNN with FPN and ResNet-50 §4.
• Teacher (EMA of the student) makes pseudo-labels from weakly augmented images; student trains on strongly augmented images §3.2.
• Pseudo-label bias from class imbalance is reduced by an EMA teacher and by a focal loss on the ROI classifier §3.3.
• Author-reported: 10.11, 11.72 and 11.60 mAP over supervised at 0.5, 1 and 2% COCO labels; 20.75 vs 13.97 for STAC at 1% tab. 1.

Problem

Semi-supervised image classification methods such as pseudo-labeling and consistency training do not transfer directly to detection. The authors argue that foreground-background imbalance and foreground-class imbalance make the pseudo-labels biased toward background and dominant classes, and that this bias is amplified when the biased pseudo-labels are added to training §1 §3.3. They also report heavy overfitting of the RPN and ROI classifiers with little labeled data, but not of box regression §1 fig. 2.

Method

The method has two stages §3 fig. 3.

1. Burn-In §3.1. Train the detector on labeled data only, with the supervised loss of eq. 1. Copy the resulting weights into the teacher and the student.
2. Teacher-Student Mutual Learning §3.2.
   • Student learning. The teacher, given a weakly augmented unlabeled image, produces pseudo-labels. The student, given a strongly augmented version of the same image, is trained on them together with the labeled loss (eq. 2). Only the student receives gradients.
   • Pseudo-label filtering. Class-wise NMS removes duplicate boxes, then boxes with confidence below δ are dropped §3.2. Box-regression losses are not applied to pseudo-labels, because confidence does not measure localization quality §3.2.
   • Teacher refinement. After each student update, the teacher weights become an exponential moving average of the student weights (eq. 3) §3.2.
   • Evaluation is done on the teacher §4.
3. Bias handling §3.3. The ROI classification loss is a multi-class focal loss in place of cross-entropy. The authors also argue that the EMA teacher limits how fast the decision boundary can move toward majority classes, and they report this effect in their ablation §4.2 fig. 5.

Equations as printed in the paper:

\mathcal{L}_{sup} = \sum_i \mathcal{L}^{rpn}_{cls}(x^s_i, y^s_i) + \mathcal{L}^{rpn}_{reg}(x^s_i, y^s_i) + \mathcal{L}^{roi}_{cls}(x^s_i, y^s_i) + \mathcal{L}^{roi}_{reg}(x^s_i, y^s_i) eq. 1

\theta_s \leftarrow \theta_s + \gamma \frac{\partial (\mathcal{L}_{sup} + \lambda_u \mathcal{L}_{unsup})}{\partial \theta_s}, \quad \mathcal{L}_{unsup} = \sum_i \mathcal{L}^{rpn}_{cls}(x^u_i, \hat{y}^u_i) + \mathcal{L}^{roi}_{cls}(x^u_i, \hat{y}^u_i) eq. 2

\theta_t \leftarrow \alpha \theta_t + (1 - \alpha) \theta_s eq. 3

Eq. 2 is printed with a plus sign in front of the gradient term. The method describes a gradient-descent update, and the code uses optimizer.step() trainer.py. Read eq. 2 as a descent step. Eq. 4 in the paper gives an unrolled form of the teacher update; see Pitfalls.

Pseudocode, following the method above and the code's order of operations:

# Burn-in: supervised loss on labeled images only  [§3.1] [eq. 1]
def supervised_loss(labeled_batch):
    # Sum of the four Faster R-CNN losses  [eq. 1]
    return rpn_cls + rpn_reg + roi_cls + roi_reg

# After burn-in: teacher starts as a copy of the student  [§3.1]
teacher_weights = student_weights

# Pseudo-labels for one unlabeled image  [§3.2] [§4]
def make_pseudo_labels(teacher, image_weak, delta):
    # Teacher sees the weakly augmented image; no gradients  [§3.2]
    detections = teacher(image_weak)
    # Class-wise NMS removes duplicate boxes  [§3.2]
    detections = class_wise_nms(detections)
    # Keep boxes whose confidence is above delta  [§3.2] [§4]
    return [d for d in detections if d.score > delta]

# Student loss for one step  [§3.2] [eq. 2]
def student_loss(labeled, unlabeled_strong, pseudo, lambda_u):
    # Supervised loss on the labeled batch  [eq. 2]
    loss_sup = supervised_loss(labeled)
    # Pseudo-label losses: RPN and ROI classification only  [§3.2] [eq. 2]
    loss_unsup = (rpn_cls_loss(unlabeled_strong, pseudo)
                  + roi_cls_loss(unlabeled_strong, pseudo))
    return loss_sup + lambda_u * loss_unsup

# Mutual-learning loop  [§3.2]
def mutual_learning(student, teacher, data, num_steps):
    for step in range(num_steps):
        labeled, unlabeled_weak, unlabeled_strong = next(data)
        # Teacher moves toward the student by EMA  [eq. 3]
        teacher.weights = (ALPHA * teacher.weights
                           + (1 - ALPHA) * student.weights)
        # Pseudo-labels come from the teacher  [§3.2]
        pseudo = make_pseudo_labels(teacher, unlabeled_weak, DELTA)
        # Student update  [eq. 2]
        loss = student_loss(labeled, unlabeled_strong, pseudo, LAMBDA_U)
        student.step(loss)

# Multi-class focal loss on the ROI classifier  [§3.3]
def focal_loss(logits, target, gamma):
    # Per-ROI softmax cross-entropy  [§3.3]
    ce = cross_entropy(logits, target)
    # Probability assigned to the true class  [§3.3]
    p = exp(-ce)
    # Easy examples (large p) are scaled down  [§3.3]
    return (1 - p) ** gamma * ce

Hyperparameters

The paper's values are in tab. 5 and app. "Training". Code values come from the COCO-standard 1% config unless stated otherwise.

Results

Author-reported. Setup: Faster R-CNN FPN, ResNet-50, Detectron2, ImageNet initialization §4. Datasets: COCO-standard (0.5 to 10% of train2017 labeled), COCO-additional (full labeled set plus COCO2017-unlabeled), VOC (VOC07 labeled, VOC12 unlabeled, test on VOC07) §4. Metric: AP50:95 §4.

COCO-standard tab. 1. Values are mean ± std. "STAC 0.5%" comes from the STAC release. CSD was re-implemented by the authors.

COCO-additional, AP50:95 tab. 2. The "3x" and "6x" labels mean 3 and 6 times 90k iterations. The text reports a 1.10 gain over supervised 3x §4.1.

VOC tab. 3. Labeled: VOC07 trainval. Supervised row is the authors' own run.

Ablations on COCO-standard 1% fig. 5. KL is the divergence between the class distribution of pseudo-labels and the ground-truth distribution.

Unsupervised loss weight λ_u on COCO-standard 10% tab. 4:

The paper's other ablations (Burn-In, threshold, EMA rate) are in the appendix app. "Additional Ablation Study" fig. 7 fig. 8 fig. 9 fig. 10.

Paper vs code

Code: official repository, commit 6977c6f, MIT licence. Paper links it in the abstract [abstract footnote].

• Burn-In length. Paper gives 1k, 2k, 6k, 12k and 20k iterations for 0.5, 1, 2, 5 and 10% labeled app. "Training". Every COCO-standard config sets BURN_UP_STEP: 2000 faster_rcnn_R_50_FPN_sup1_run1.yaml (same line in the 0.5, 2, 5 and 10% configs).
• VOC hyperparameters. Paper says λ_u = 4 and b_l = b_u = 32 for VOC tab. 5. The VOC config uses λ_u = 2.0, EMA 0.999 and batch 16 voc07_voc12.yaml voc07_voc12.yaml.
• COCO-additional batch. Paper's Table 5 says b_l = b_u = 16. The config and the appendix text say 32 faster_rcnn_R_50_FPN_focal_cross_coco_unsupW2.yaml.
• Learning-rate schedule. Paper says constant LR app. "Training". COCO-standard configs set STEPS: (179990, 179995) faster_rcnn_R_50_FPN_sup1_run1.yaml. The configs do not set GAMMA or warmup, so Detectron2 defaults apply. I did not verify those defaults. COCO-additional uses FACTOR_LIST: (1, 1, 1, 1, 1), which is constant faster_rcnn_R_50_FPN_focal_cross_coco_unsupW2.yaml.
• Labeled images are used with two views. Burn-In and mutual learning both feed the labeled batch twice, once from each augmentation trainer.py trainer.py. Eq. 1 and eq. 2 do not state this.
• Which targets train the RPN. The code thresholds RPN proposals too trainer.py, but the pseudo-labels added to the training data are the ROI-head outputs trainer.py. The RPN classification loss therefore uses ROI-head pseudo-labels. Paper's ŷ is not defined per head.
• Confidence threshold on the RPN. The code thresholds RPN objectness_logits directly against δ = 0.7 trainer.py. Whether these are probabilities or raw logits depends on Detectron2, which is not in this repository. I did not verify this.
• Class-wise NMS. Paper describes NMS before thresholding, following STAC §3.2. The ROI pseudo-labels come from Detectron2's box-predictor inference roi_heads.py. The repository does not set an NMS threshold, so the value used is Detectron2's default. I did not verify that default.
• Box regression. Consistent with the paper: pseudo-label losses for box regression are multiplied by 0 trainer.py.
• Unsupervised weight. Applied to each pseudo-label loss trainer.py.
• Strong augmentation. Matches tab. 6 detection_utils.py.
• Teacher update timing. Code copies the student into the teacher at Burn-In end, then updates by EMA at the start of each later iteration, before pseudo-labeling trainer.py. The paper gives no order.
• Teacher in train mode. Code runs the teacher in training mode under no_grad, so batch-norm statistics are updated by the forward pass trainer.py. The EMA also covers the whole state_dict trainer.py. The paper does not mention either.
• Evaluation. Paper reports the teacher §4. The code evaluates both student and teacher trainer.py.
• Focal loss. Code uses gamma = 1.5 and normalises by the number of sampled ROIs fast_rcnn.py. The paper does not state gamma.

Pitfalls

• Equation 2 sign. The printed update is θ_s + γ ∂L/∂θ_s. Use descent (minus), matching the code's optimizer step.
• Equation 4. The paper gives an unrolled sum for the teacher weights, with γ, α and an exponent that depends on k and i. I could not reconcile it with eq. 3, and the code only implements eq. 3 recursively. Do not implement eq. 4 without checking it yourself.
• α is the teacher's share. In eq. 3, α = 0.9996 means the teacher keeps 99.96% of its weights each step. Using 0.9996 as the student's share would invert the method. The α ablation text says the best value is 0.99 app. "Additional Ablation Study", yet every run uses 0.9996. Treat the ablation as a sweep, not a tuned value.
• Ablation vs headline. The ablation's Focal + EMA result on COCO 1% is 21.19 mAP fig. 5. The main table reports 20.75 for the same setting tab. 1. The paper does not explain the gap.
• λ_u ablation. The best result in the sweep is λ_u = 5.0 (32.00), but the method uses λ_u = 4.0 tab. 4. λ_u = 8.0 does not converge.
• Threshold stability. The appendix reports that thresholds below 0.5 do not converge and that δ = 0.9 gives too few pseudo-labels [app. "Additional Ablation Study"].
• Burn-In sensitivity. The appendix reports that Burn-In gives more accurate early pseudo-boxes and better final results app. "Additional Ablation Study" fig. 7. Burn-In length is a per-run setting in the code, so use the config value for each label fraction.
• Compute. COCO-standard runs are 180k iterations, each with a batch of 32 labeled and 32 unlabeled images, plus a teacher forward pass app. "Training". The paper does not report wall-clock time or hardware in the sections read.
• Figure and table numbers. Figure and table numbers in this card follow the order of their captions in the LaTeX source. Check them against the compiled PDF before quoting them.
• Terminology. The paper's γ is the learning rate in eq. 2 and the focal loss exponent in §3.3 is also commonly written γ. Do not mix the two.
• Appendix naming. The appendix calls the COCO-additional Burn-In stage "Burn-Up". The config name BURN_UP_STEP is the same concept.

Lineage

• Builds on: STAC, a pseudo-labeling detector with a fixed teacher (arXiv:2005.04757), which the paper compares against §4.
• Also builds on: CSD, consistency-based semi-supervised detection (paper's reference "jeong2019consistency") §2; FixMatch, the pseudo-labeling classifier with weak and strong views §2; Mean Teacher, the EMA teacher §3.2.
• Uses focal loss for the ROI classifier §3.3 [lin2017focal].
• Compared against: CSD and STAC on COCO-standard and VOC tab. 1 tab. 3.
• Superseded by: not known from content in this wiki.