基于软教师的端到端半监督目标检测

作者：徐孟德*、张政*、胡汉、王建峰、王丽娟、魏方云、白翔、刘子成。

本仓库是 ICCV2021 论文《基于软教师的端到端半监督目标检测》（https://arxiv.org/abs/2106.09018）的官方实现。

引用

@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}
}

主要结果

部分标注数据

我们遵循 STAC[1] 的方法，针对每种设置评估了 5 种不同的数据划分，并报告了 5 次划分的平均性能。结果如下：

1% 标注数据

方法	mAP	模型权重	配置文件
基线	10.0	-	配置
我们方法 (thr=5e-2)	21.62	网盘	配置
我们方法 (thr=1e-3)	22.64	网盘	配置

5% 标注数据

方法	mAP	模型权重	配置文件
基线	20.92	-	配置
我们方法 (thr=5e-2)	30.42	网盘	配置
我们方法 (thr=1e-3)	31.7	网盘	配置

10% 标注数据

方法	mAP	模型权重	配置文件
基线	26.94	-	配置
我们方法 (thr=5e-2)	33.78	网盘	配置
我们方法 (thr=1e-3)	34.7	网盘	配置

全部标注数据

Faster R-CNN (ResNet-50)

模型	mAP	模型权重	配置文件
基线	40.9	-	配置
我们方法 (thr=5e-2)	44.05	网盘	配置
我们方法 (thr=1e-3)	44.6	网盘	配置
我们方法* (thr=5e-2)	44.5	-	配置
我们方法* (thr=1e-3)	44.9	-	配置

Faster R-CNN (ResNet-101)

模型	mAP	模型权重	配置文件
基线	43.8	-	配置
我们方法* (thr=5e-2)	46.9	网盘	配置
我们方法* (thr=1e-3)	47.6	网盘	配置

注释

• 我们的*表示我们使用了更长的训练计划。
• thr 表示配置文件中的 model.test_cfg.rcnn.score_thr。这一推理技巧最早由 Instant-Teaching[2] 提出。
• 所有模型均在 8 张 V100 GPU 上进行训练。

使用方法

需求

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

注意事项

• 我们使用 wandb 进行可视化，如果您不想使用它，只需注释掉 configs/soft_teacher/base.py 文件中的第 273–284 行即可。
• 该项目应与最新版本的 mmdetection 兼容。如果您想切换到与我们相同的 mmdetection 版本，请运行 cd thirdparty/mmdetection && git checkout v2.16.0。

安装

make install

数据准备

• 下载 COCO 数据集
• 执行以下命令以生成数据集划分：

# YOUR_DATA 应该是一个包含 COCO 数据集的目录。
# 例如：
# YOUR_DATA/
#  coco/
#     train2017/
#     val2017/
#     unlabeled2017/
#     annotations/
ln -s ${YOUR_DATA} data
bash tools/dataset/prepare_coco_data.sh conduct

有关具体需要下载的内容，请参阅 tools/dataset/prepare_coco_data.sh 文件中的 11-24 行。

训练

• 在部分标注数据设置下训练模型：

# JOB_TYPE: 'baseline' 或 'semi'，决定运行哪种任务
# PERCENT_LABELED_DATA: 1、5、10，表示标注的 COCO 数据占整个训练数据集的比例。
# GPU_NUM: 用于运行任务的 GPU 数量
for FOLD in 1 2 3 4 5;
do
  bash tools/dist_train_partially.sh <JOB_TYPE> ${FOLD} <PERCENT_LABELED_DATA> <GPU_NUM>
done

例如，我们可以运行以下脚本，使用 8 张 GPU 在 10% 标注数据上训练我们的模型：

for FOLD in 1 2 3 4 5;
do
  bash tools/dist_train_partially.sh semi ${FOLD} 10 8
done

• 在全标注数据设置下训练模型：

bash tools/dist_train.sh <CONFIG_FILE_PATH> <NUM_GPUS>

例如，使用 8 张 GPU 训练我们的 R50 模型：

bash tools/dist_train.sh configs/soft_teacher/soft_teacher_faster_rcnn_r50_caffe_fpn_coco_full_720k.py 8

• 在新数据集上训练模型：

核心思想是将新数据集转换为 COCO 格式。详细信息请参阅 添加新数据集。

评估

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

推理

使用训练好的模型进行推理并可视化检测结果：

# [IMAGE_FILE_PATH]: 您本地文件系统中图像文件的路径
# [CONFIG_FILE]: 配置文件的路径
# [CHECKPOINT_PATH]: 与所提供配置文件相关的已训练模型的路径
# [OUTPUT_PATH]: 保存检测结果的目录
python demo/image_demo.py [IMAGE_FILE_PATH] [CONFIG_FILE] [CHECKPOINT_PATH] --output [OUTPUT_PATH]

例如：

• 使用提供的 R50 模型对单张图像进行推理：

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/

程序完成后，会生成一张与输入图像同名的图片，并保存到 work_dirs 目录中。

• 使用提供的 R50 模型对多张图像进行推理：

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] 一种用于目标检测的简单半监督学习框架

[2] Instant-Teaching：一个端到端的半监督目标检测框架

SoftTeacher 快速上手指南

SoftTeacher 是一个基于 MMDetection 的端到端半监督目标检测框架，源自 ICCV 2021 论文。它通过“软教师”机制利用未标记数据显著提升检测性能，尤其适用于标注数据稀缺的场景。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Ubuntu 16.04 或更高版本
• Python: 3.6 (推荐通过 Anaconda3 管理)
• PyTorch: 1.9.0
• 核心依赖:
  • mmdetection == 2.16.0+fe46ffe
  • mmcv-full == 1.3.9
  • wandb == 0.10.31 (用于可视化，可选)

注意：虽然项目理论上兼容最新版 MMDetection，但为了复现论文结果，建议严格匹配上述版本。国内用户可使用清华源或阿里源加速 pip 和 conda 安装。

安装步骤

1. 克隆仓库并安装依赖 执行以下命令自动完成环境配置和依赖安装：

   make install
   

2. 验证 MMDetection 版本（可选） 如果需要强制切换到与本项目完全一致的 MMDetection 版本：

   cd thirdparty/mmdetection && git checkout v2.16.0
   

3. 配置 WandB（可选） 如果您不需要使用 WandB 进行训练日志可视化，可以编辑 configs/soft_teacher/base.py 文件，注释掉第 273-284 行。

基本使用

1. 数据准备

本项目基于 COCO 数据集。请下载 COCO 数据集（包含 train2017, val2017, annotations 等），并按以下结构组织目录，然后生成半监督学习所需的数据划分：

# 假设 ${YOUR_DATA} 是包含 coco 数据集的根目录
# 例如：${YOUR_DATA}/coco/train2017/, ${YOUR_DATA}/coco/annotations/ 等
ln -s ${YOUR_DATA} data

# 生成数据划分文件
bash tools/dataset/prepare_coco_data.sh conduct

2. 模型训练

场景 A：半监督训练（部分标注数据）

适用于只有 1%、5% 或 10% 标注数据的场景。以下示例展示如何使用 8 张 GPU 在 10% 标注数据上训练（遍历 5 个折叠 fold）：

for FOLD in 1 2 3 4 5;
do
  bash tools/dist_train_partially.sh semi ${FOLD} 10 8
done

• semi: 任务类型（半监督）。
• 10: 标注数据比例（可选 1, 5, 10）。
• 8: 使用的 GPU 数量。

场景 B：全监督训练（全量标注数据）

使用全量 COCO 数据进行训练：

bash tools/dist_train.sh configs/soft_teacher/soft_teacher_faster_rcnn_r50_caffe_fpn_coco_full_720k.py 8

3. 模型推理与可视化

使用训练好的模型对单张或多张图片进行推理，并保存检测结果：

python demo/image_demo.py [IMAGE_FILE_PATH] [CONFIG_FILE] [CHECKPOINT_PATH] --output [OUTPUT_PATH]

示例：对单张图片进行推理

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/

示例：对目录下所有图片进行推理

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/

运行完成后，带有检测框的图片将保存在指定的输出目录中。