高质量、快速、模块化的 SSD PyTorch 1.0 参考实现

本仓库实现了 SSD（单次多框检测器）。该实现深受 ssd.pytorch、pytorch-ssd 和 maskrcnn-benchmark 等项目的影响。本仓库旨在成为基于 SSD 的研究代码基础。

损失函数	学习率	指标

亮点

• PyTorch 1.0：支持 PyTorch 1.0 或更高版本。
• 多 GPU 训练与推理：我们使用 DistributedDataParallel，您可以使用任意数量的 GPU 进行训练或测试，训练方案会相应调整。
• 模块化：轻松添加自定义模块。我们抽象了 backbone、Detector、BoxHead、BoxPredictor 等组件，您可以替换其中任何部分而无需修改现有代码。例如，您可以将 EfficientNet 作为骨干网络，只需添加 efficient_net.py（已包含）并注册，在配置文件中指定即可，完成！
• 推理时支持 CPU：在推理阶段可在 CPU 上运行。
• 流畅愉悦的训练流程：我们会保存模型、优化器、调度器和训练迭代的状态，您可以在任意时刻停止训练，并从保存点继续，无需更改训练命令。
• 批处理推理：每个 GPU 每批次可同时对多张图像进行推理。
• 训练期间评估：每 eval_step 对模型进行评估，以检查性能是否提升。
• 指标可视化：在 TensorBoard 中可视化指标详情，如 COCO 数据集的 AP、APl、APm 和 APs，或 VOC 数据集的 mAP 和 20 类别的 AP。
• 自动下载：从 URL 加载预训练权重并缓存。

安装

要求

1. Python3
2. PyTorch 1.0 或更高版本
3. yacs
4. Vizer
5. GCC >= 4.9
6. OpenCV

安装步骤

git clone https://github.com/lufficc/SSD.git
cd SSD
# 必需包：torch torchvision yacs tqdm opencv-python vizer
pip install -r requirements.txt

# 完成！就这么简单！无需编译！无需繁琐的安装过程！

# 建议安装最新版本的 torch 和 torchvision。

训练

数据集设置

Pascal VOC

对于 Pascal VOC 数据集，文件夹结构应如下所示：

VOC_ROOT
|__ VOC2007
    |_ JPEGImages
    |_ Annotations
    |_ ImageSets
    |_ SegmentationClass
|__ VOC2012
    |_ JPEGImages
    |_ Annotations
    |_ ImageSets
    |_ SegmentationClass
|__ ...

其中 VOC_ROOT 默认为当前项目的 datasets 文件夹，您也可以创建指向 datasets 的符号链接，或通过 export VOC_ROOT="/path/to/voc_root" 来指定路径。

COCO

对于 COCO 数据集，文件夹结构应如下所示：

COCO_ROOT
|__ annotations
    |_ instances_valminusminival2014.json
    |_ instances_minival2014.json
    |_ instances_train2014.json
    |_ instances_val2014.json
    |_ ...
|__ train2014
    |_ <im-1-name>.jpg
    |_ ...
    |_ <im-N-name>.jpg
|__ val2014
    |_ <im-1-name>.jpg
    |_ ...
    |_ <im-N-name>.jpg
|__ ...

其中 COCO_ROOT 默认为当前项目的 datasets 文件夹，您也可以创建指向 datasets 的符号链接，或通过 export COCO_ROOT="/path/to/coco_root" 来指定路径。

单 GPU 训练

# 例如，训练 SSD300：
python train.py --config-file configs/vgg_ssd300_voc0712.yaml

多 GPU 训练

# 例如，使用 4 个 GPU 训练 SSD300：
export NGPUS=4
python -m torch.distributed.launch --nproc_per_node=$NGPUS train.py --config-file configs/vgg_ssd300_voc0712.yaml SOLVER.WARMUP_FACTOR 0.03333 SOLVER.WARMUP_ITERS 1000

我提供的配置文件假设是在单 GPU 上运行。当 GPU 数量发生变化时，超参数（学习率、最大迭代次数等）也需要根据以下论文进行调整：准确的大批量 SGD：1 小时内训练 ImageNet。

评估

单 GPU 评估

# 例如，评估 SSD300：
python test.py --config-file configs/vgg_ssd300_voc0712.yaml

多 GPU 评估

# 例如，使用 4 个 GPU 评估 SSD300：
export NGPUS=4
python -m torch.distributed.launch --nproc_per_node=$NGPUS test.py --config-file configs/vgg_ssd300_voc0712.yaml

示例

对文件夹中的图像进行预测非常简单：

python demo.py --config-file configs/vgg_ssd300_voc0712.yaml --images_dir demo --ckpt https://github.com/lufficc/SSD/releases/download/1.2/vgg_ssd300_voc0712.pth

程序会自动下载并缓存 vgg_ssd300_voc0712.pth，预测结果（包含边界框、置信度分数和类别名称）默认会保存到 demo/result 文件夹中。

您将看到类似以下的输出：

(0001/0005) 004101.jpg: 检测到 1 个目标 | 加载耗时 010ms | 推理耗时 033ms | FPS 31
(0002/0005) 003123.jpg: 检测到 5 个目标 | 加载耗时 009ms | 推理耗时 019ms | FPS 53
(0003/0005) 000342.jpg: 检测到 2 个目标 | 加载耗时 009ms | 推理耗时 019ms | FPS 51
(0004/0005) 008591.jpg: 检测到 2 个目标 | 加载耗时 008ms | 推理耗时 020ms | FPS 50
(0005/0005) 000542.jpg: 检测到 1 个目标 | 加载耗时 011ms | 推理耗时 019ms | FPS 53

模型库

原始论文：

	VOC2007 测试	coco test-dev2015
SSD300*	77.2	25.1
SSD512*	79.8	28.8

COCO：

骨干网络	输入尺寸	box AP	模型大小	下载
VGG16	300	25.2	262MB	模型
VGG16	512	29.0	275MB	模型

PASCAL VOC：

主干网络	输入尺寸	mAP	模型大小	下载
VGG16	300	77.7	201MB	模型
VGG16	512	80.7	207MB	模型
Mobilenet V2	320	68.9	25.5MB	模型
Mobilenet V3	320	69.5	29.9MB	模型
EfficientNet-B3	300	73.9	97.1MB	模型

开发指南

如果您想添加自定义组件，请参阅 DEVELOP_GUIDE.md，以获取更多详细信息。

故障排除

如果您在运行或编译此代码时遇到问题，我们已在 TROUBLESHOOTING.md 中整理了一份常见问题列表。如果您的问题未在此列出，请随时提交一个新的问题。

引用

如果您在研究中使用了本项目，请引用该项目。

@misc{lufficc2018ssd,
    author = {Congcong Li},
    title = {{PyTorch 中高质量、快速、模块化的 SSD 参考实现}},
    year = {2018},
    howpublished = {\url{https://github.com/lufficc/SSD}}
}

SSD 快速上手指南

本指南基于 lufficc/SSD 项目，这是一个高质量、快速且模块化的 PyTorch 1.0+ SSD（Single Shot MultiBox Detector）参考实现。

环境准备

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

• 操作系统: Linux (推荐)
• Python: Python 3
• 深度学习框架: PyTorch 1.0 或更高版本
• 编译器: GCC >= 4.9
• 其他依赖:
  • yacs (配置管理)
  • opencv-python (图像处理)
  • vizer (可视化工具，需从 GitHub 安装)
  • tqdm (进度条)

提示：国内用户建议使用清华源或阿里源加速 pip 包的安装。

安装步骤

本项目无需复杂的编译过程（No BUILD），只需克隆代码并安装依赖即可。

1. 克隆仓库

   git clone https://github.com/lufficc/SSD.git
   cd SSD
   

2. 安装依赖 建议先安装最新版的 torch 和 torchvision，然后安装项目所需的其他包。

   通用安装命令：

   pip install -r requirements.txt
   

   国内加速安装示例（使用清华源）：

   pip install torch torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple
   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   注意：vizer 库如果在国内源找不到，可能需要通过 pip install git+https://github.com/lufficc/Vizer.git 进行安装。

安装完成后，无需额外设置即可直接使用。

基本使用

1. 快速演示 (Demo)

最简单的使用方式是直接运行预训练模型对图片进行检测。以下命令会自动下载预训练权重并进行推理：

python demo.py --config-file configs/vgg_ssd300_voc0712.yaml --images_dir demo --ckpt https://github.com/lufficc/SSD/releases/download/1.2/vgg_ssd300_voc0712.pth

• 功能：自动下载 vgg_ssd300_voc0712.pth 权重文件并缓存。
• 输出：检测结果（包含边界框、分数和标签）将默认保存在 demo/result 文件夹中。
• 终端反馈：运行时会显示每张图的加载时间、推理时间和 FPS。

2. 数据集准备 (以 Pascal VOC 为例)

若要训练或评估自己的模型，需按以下结构组织数据（默认为项目下的 datasets 文件夹，也可通过环境变量 VOC_ROOT 指定）：

VOC_ROOT
|__ VOC2007
    |_ JPEGImages
    |_ Annotations
    |_ ImageSets
|__ VOC2012
    |_ JPEGImages
    |_ Annotations
    |_ ImageSets

3. 开始训练

单卡训练示例：

python train.py --config-file configs/vgg_ssd300_voc0712.yaml

多卡训练示例 (4 张 GPU)：

export NGPUS=4
python -m torch.distributed.launch --nproc_per_node=$NGPUS train.py --config-file configs/vgg_ssd300_voc0712.yaml SOLVER.WARMUP_FACTOR 0.03333 SOLVER.WARMUP_ITERS 1000

4. 模型评估

单卡评估：

python test.py --config-file configs/vgg_ssd300_voc0712.yaml

多卡评估：

export NGPUS=4
python -m torch.distributed.launch --nproc_per_node=$NGPUS test.py --config-file configs/vgg_ssd300_voc0712.yaml

---

更多自定义模块开发指南请参考项目中的 DEVELOP_GUIDE.md。