SSD：单次多框目标检测器，基于 PyTorch

这是 Wei Liu、Dragomir Anguelov、Dumitru Erhan、Christian Szegedy、Scott Reed、Cheng-Yang 和 Alexander C. Berg 在 2016 年发表的论文中提出的 Single Shot MultiBox Detector 的 PyTorch 实现。官方原始的 Caffe 代码可以在 这里找到。

目录

• 安装
• 数据集
• 训练
• 评估
• 性能
• 演示
• 未来工作
• 参考文献

       

安装

• 根据你的环境在 PyTorch 官网选择并运行相应的命令来安装 PyTorch。
• 克隆本仓库。
  • 注意：我们目前仅支持 Python 3 及以上版本。
• 然后按照下面的说明下载数据集。
• 我们现在支持使用 Visdom 在训练过程中进行实时损失可视化！
  • 要在浏览器中使用 Visdom：

  # 首先安装 Python 服务器和客户端
  pip install visdom
  # 启动服务器（最好在 screen 或 tmux 中）
  python -m visdom.server
  

  • 然后（在训练期间）访问 http://localhost:8097/（训练细节请参见下方的“训练”部分）。
• 注意：对于训练，我们目前支持 VOC 和 COCO 数据集，并计划尽快添加对 ImageNet 的支持。

数据集

为了方便起见，我们提供了 bash 脚本来帮你处理数据集的下载和设置。我们还提供了简单的数据集加载器，它们继承自 torch.utils.data.Dataset，因此与 torchvision.datasets 的 API 完全兼容。

COCO

Microsoft COCO：上下文中的常见物体

下载 COCO 2014

# 指定数据集下载的目标目录，否则默认为 ~/data/
sh data/scripts/COCO2014.sh

VOC 数据集

PASCAL VOC：视觉对象类别

下载 VOC2007 trainval 和 test

# 指定数据集下载的目标目录，否则默认为 ~/data/
sh data/scripts/VOC2007.sh # <directory>

下载 VOC2012 trainval

# 指定数据集下载的目标目录，否则默认为 ~/data/
sh data/scripts/VOC2012.sh # <directory>

训练 SSD

• 首先下载 fc-reduced VGG-16 的 PyTorch 基础网络权重，地址为： https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth
• 默认情况下，我们认为你已将该文件下载到 ssd.pytorch/weights 目录：

mkdir weights
cd weights
wget https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth

• 要使用训练脚本训练 SSD，只需在 train.py 中列出的参数中指定标志或手动更改它们即可。

python train.py

• 注意：
  • 对于训练，强烈建议使用 NVIDIA GPU 以提高速度。
  • 关于 Visdom 的使用和安装说明，请参阅< a href='#installation'>安装部分。
  • 你可以通过指定检查点路径从检查点继续训练（同样，请参阅 train.py 中的选项）

评估

要评估训练好的网络：

python eval.py

你可以通过标记或手动更改 eval.py 文件中列出的参数来指定它们。

性能

VOC2007 测试

mAP

原始	转换后的 weiliu89 权重	从零开始无数据增强	从零开始有数据增强
77.2 %	77.26 %	58.12%	77.43 %

FPS

GTX 1060: ~45.45 FPS

演示

使用预训练的 SSD 网络进行检测

下载预训练的网络

• 我们正努力提供在不同数据集上训练的最新 SSD 模型定义的 PyTorch state_dicts（权重张量字典）。
• 当前，我们提供以下 PyTorch 模型：
  • 在 VOC0712 上训练的 SSD300（最新的 PyTorch 权重）
    • https://s3.amazonaws.com/amdegroot-models/ssd300_mAP_77.43_v2.pth
  • 在 VOC0712 上训练的 SSD300（原始 Caffe 权重）
    • https://s3.amazonaws.com/amdegroot-models/ssd_300_VOC0712.pth
• 我们的目标是重现原始论文中的这张表

尝试演示笔记本

• 确保你已安装 jupyter notebook。
• 安装 jupyter notebook 的两种方法：

  1. 如果你使用 conda 安装了 PyTorch（推荐），那么你应该已经拥有它了。只需导航到克隆的 ssd.pytorch 仓库并运行： jupyter notebook

  2. 如果使用 pip：

# 确保 pip 已升级
pip3 install --upgrade pip
# 安装 jupyter notebook
pip install jupyter
# 在 ssd.pytorch 内运行
jupyter notebook

• 现在导航到 http://localhost:8888（默认）的 demo/demo.ipynb，尽情体验吧！

尝试摄像头演示

• 可以在 CPU 上运行（可能需要调整 cv2.waitkey 以获得最佳帧率）或在 NVIDIA GPU 上运行
• 此演示目前需要 opencv2+ 和 Python 绑定以及内置摄像头
  • 你可以在 demo/live.py 中更改默认摄像头
• 安装 imutils 包以利用 CPU 上的多线程功能：
  • pip install imutils
• 运行 python -m demo.live 即可打开摄像头并开始检测！

TODO

我们整理了一份待办事项清单，希望在不久的将来完成：

• 接下来还有：
  • 支持 MS COCO 数据集
  • 支持 SSD512 的训练和测试
  • 支持在自定义数据集上进行训练

作者

• Max deGroot
• Ellis Brown

注： 很遗憾，这仅仅是我们的一项业余爱好，而不是全职工作，因此我们会尽力保持内容的更新，但无法保证。尽管如此，感谢大家一直以来的帮助和反馈，我们非常感激。我们会尽快处理所有问题。

参考文献

• 魏 Liu 等人. “SSD：单次多框检测器.” ECCV2016.
• 原始实现（CAFFE）
• 衷心感谢 Alex Koltun 及其在 Webyclip 的团队，在完成数据增强部分时给予的帮助。
• 其他优秀的 SSD 移植项目列表，这些项目曾为我们提供了灵感（尤其是 Chainer 仓库）：
  • Chainer, Keras, MXNet, TensorFlow

SSD.pytorch 快速上手指南

SSD (Single Shot MultiBox Detector) 是一个基于 PyTorch 实现的单阶段目标检测算法。本指南将帮助你快速搭建环境并运行模型。

1. 环境准备

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

• 操作系统: Linux 或 macOS (Windows 支持需自行配置)
• Python: 3.0 及以上版本
• 深度学习框架: PyTorch (建议安装最新稳定版)
• 硬件: 训练过程强烈推荐使用 NVIDIA GPU；推理可在 CPU 上进行（速度较慢）
• 可选依赖:
  • visdom: 用于训练过程中的实时损失可视化
  • opencv-python: 用于演示和视频流检测
  • jupyter: 用于运行示例 Notebook

安装基础依赖命令：

# 安装 PyTorch (请访问 pytorch.org 获取适合你环境的命令，以下为示例)
pip install torch torchvision

# 安装可视化工具和图像处理库
pip install visdom opencv-python imutils jupyter

提示：国内用户建议使用清华源或阿里源加速 pip 安装： pip install -i https://pypi.tuna.tsinghua.edu.cn/simple <package_name>

2. 安装步骤

2.1 克隆项目

首先将代码仓库克隆到本地：

git clone https://github.com/amdegroot/ssd.pytorch.git
cd ssd.pytorch

2.2 下载预训练骨干网络

训练前需要下载经过裁剪的 VGG-16 预训练权重作为骨干网络。执行以下命令自动下载至 weights 目录：

mkdir weights
cd weights
wget https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth
cd ..

注意：如果 wget 下载速度慢，可手动在浏览器下载该文件并放入 ssd.pytorch/weights/ 目录下。

2.3 数据集准备 (可选)

如果你打算重新训练模型，需要下载数据集。项目提供了自动化脚本（默认下载到 ~/data/ 目录）：

• COCO 2014:

  sh data/scripts/COCO2014.sh
  

• PASCAL VOC 2007 & 2012:

  sh data/scripts/VOC2007.sh
  sh data/scripts/VOC2012.sh
  

3. 基本使用

3.1 运行演示 (最快体验)

无需训练，直接下载预训练模型进行物体检测演示。

1. 下载预训练模型 (以 VOC 数据集训练的 SSD300 为例)：

   mkdir weights
   cd weights
   wget https://s3.amazonaws.com/amdegroot-models/ssd300_mAP_77.43_v2.pth
   cd ..
   

2. 启动 Jupyter Notebook 查看示例：

   jupyter notebook
   

   在浏览器中打开 demo/demo.ipynb，按照单元格顺序运行即可看到检测效果。

3. 或者运行摄像头实时检测 (需连接摄像头)：

   python -m demo.live
   

3.2 训练模型

使用默认参数开始训练（需先准备好数据集和 VGG 权重）：

python train.py

• 可视化监控：在另一个终端启动 Visdom 服务，然后在浏览器访问 http://localhost:8097/ 查看实时训练曲线。

  python -m visdom.server
  

3.3 评估模型

对训练好的网络进行评估：

python eval.py

你可以通过修改 train.py 或 eval.py 中的参数，或在命令行中添加标志位来调整超参数和数据集路径。