Skip to content

Repository files navigation

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

PWCPWCPWCPWCPWCPWCPWCPWC

By Mengde Xu*, Zheng Zhang*, Han Hu, Jianfeng Wang, Lijuan Wang, Fangyun Wei, Xiang Bai, Zicheng Liu.

This repo is the official implementation of ICCV2021 paper "End-to-End Semi-Supervised Object Detection with Soft Teacher".

Citation

@article{xu2021end,
title={End-to-End Semi-Supervised Object Detection with Soft Teacher},
author={Xu, Mengde and Zhang, Zheng and Hu, Han and Wang, Jianfeng and Wang, Lijuan and Wei, Fangyun and Bai, Xiang and Liu, Zicheng},
journal={Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
year={2021}
}

Main Results

Partial Labeled Data

We followed STAC[1] to evaluate on 5 different data splits for each setting, and report the average performance of 5 splits. The results are shown in the following:

1% labeled data

MethodmAPModel WeightsConfig Files
Baseline10.0-Config
Ours (thr=5e-2)21.62DriveConfig
Ours (thr=1e-3)22.64DriveConfig

5% labeled data

MethodmAPModel WeightsConfig Files
Baseline20.92-Config
Ours (thr=5e-2)30.42DriveConfig
Ours (thr=1e-3)31.7DriveConfig

10% labeled data

MethodmAPModel WeightsConfig Files
Baseline26.94-Config
Ours (thr=5e-2)33.78DriveConfig
Ours (thr=1e-3)34.7DriveConfig

Full Labeled Data

Faster R-CNN (ResNet-50)

ModelmAPModel WeightsConfig Files
Baseline40.9-Config
Ours (thr=5e-2)44.05DriveConfig
Ours (thr=1e-3)44.6DriveConfig
Ours* (thr=5e-2)44.5-Config
Ours* (thr=1e-3)44.9-Config

Faster R-CNN (ResNet-101)

ModelmAPModel WeightsConfig Files
Baseline43.8-Config
Ours* (thr=5e-2)46.9DriveConfig
Ours* (thr=1e-3)47.6DriveConfig

Notes

  • Ours* means we use longer training schedule.
  • thr indicates model.test_cfg.rcnn.score_thr in config files. This inference trick was first introduced by Instant-Teaching[2].
  • All models are trained on 8*V100 GPUs

Usage

Requirements

  • Ubuntu 16.04
  • Anaconda3 with python=3.6
  • Pytorch=1.9.0
  • mmdetection=2.16.0+fe46ffe
  • mmcv=1.3.9
  • wandb=0.10.31

Notes

  • We use wandb for visualization, if you don't want to use it, just comment line 273-284 in configs/soft_teacher/base.py.
  • The project should be compatible to the latest version of mmdetection. If you want to switch to the same version mmdetection as ours, run cd thirdparty/mmdetection && git checkout v2.16.0

Installation

make install

Data Preparation

  • Download the COCO dataset
  • Execute the following command to generate data set splits:
# YOUR_DATA should be a directory contains coco dataset.# For eg.:# YOUR_DATA/# coco/# train2017/# val2017/# unlabeled2017/# annotations/
ln -s ${YOUR_DATA} data
bash tools/dataset/prepare_coco_data.sh conduct

For concrete instructions of what should be downloaded, please refer to tools/dataset/prepare_coco_data.sh line 11-24

Training

  • To train model on the partial labeled data setting:
# JOB_TYPE: 'baseline' or 'semi', decide which kind of job to run# PERCENT_LABELED_DATA: 1, 5, 10. The ratio of labeled coco data in whole training dataset.# GPU_NUM: number of gpus to run the jobforFOLDin 1 2 3 4 5;do
bash tools/dist_train_partially.sh <JOB_TYPE>${FOLD}<PERCENT_LABELED_DATA><GPU_NUM>done

For example, we could run the following scripts to train our model on 10% labeled data with 8 GPUs:

forFOLDin 1 2 3 4 5;do
bash tools/dist_train_partially.sh semi ${FOLD} 10 8
done
  • To train model on the full labeled data setting:
bash tools/dist_train.sh <CONFIG_FILE_PATH><NUM_GPUS>

For example, to train ours R50 model with 8 GPUs:

bash tools/dist_train.sh configs/soft_teacher/soft_teacher_faster_rcnn_r50_caffe_fpn_coco_full_720k.py 8
  • To train model on new dataset:

The core idea is to convert a new dataset to coco format. Details about it can be found in the adding new dataset.

Evaluation

bash tools/dist_test.sh <CONFIG_FILE_PATH> <CHECKPOINT_PATH> <NUM_GPUS> --eval bbox --cfg-options model.test_cfg.rcnn.score_thr=<THR>

Inference

To inference with trained model and visualize the detection results:

# [IMAGE_FILE_PATH]: the path of your image file in local file system# [CONFIG_FILE]: the path of a confile file# [CHECKPOINT_PATH]: the path of a trained model related to provided confilg file.# [OUTPUT_PATH]: the directory to save detection result
python demo/image_demo.py [IMAGE_FILE_PATH] [CONFIG_FILE] [CHECKPOINT_PATH] --output [OUTPUT_PATH]

For example:

  • Inference on single image with provided R50 model:
python demo/image_demo.py /tmp/tmp.png configs/soft_teacher/soft_teacher_faster_rcnn_r50_caffe_fpn_coco_full_720k.py work_dirs/downloaded.model --output work_dirs/

After the program completes, a image with the same name as input will be saved to work_dirs

  • Inference on many images with provided R50 model:
python demo/image_demo.py '/tmp/*.jpg' configs/soft_teacher/soft_teacher_faster_rcnn_r50_caffe_fpn_coco_full_720k.py work_dirs/downloaded.model --output work_dirs/

[1] A Simple Semi-Supervised Learning Framework for Object Detection

[2] Instant-Teaching: An End-to-End Semi-Supervised Object Detection Framework

About

Semi-Supervised Learning, Object Detection, ICCV2021

Topics

Resources

Code of conduct

Security policy

Stars

928 stars

Watchers

22 watching

Forks

Used by

Contributors

Languages