SSD:单次多框目标检测器,基于 PyTorch
PyTorch 实现的单次多框检测器,源自 Wei Liu, Dragomir Anguelov, Dumitru Erhan, Christian Szegedy, Scott Reed, Cheng-Yang 和 Alexander C. Berg 于 2016 年发表的论文。官方的原始 Caffe 代码可以在这里找到。
目录
安装
- 通过在官网上选择您的环境并运行相应的命令,安装 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
微软 COCO:上下文中的常见物体
下载 COCO 2014
# 指定一个目录用于下载数据集,否则默认为 ~/data/
sh data/scripts/COCO2014.sh
VOC 数据集
PASCAL VOC:视觉对象类别
下载 VOC2007 训练集和测试集
# 指定一个目录用于下载数据集,否则默认为 ~/data/
sh data/scripts/VOC2007.sh # <directory>
下载 VOC2012 训练集
# 指定一个目录用于下载数据集,否则默认为 ~/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 使用/安装的说明,请参见安装部分。
- 您可以通过指定路径作为训练参数之一(再次参见
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 权重)
- 在 VOC0712 上训练的 SSD300(原始 Caffe 权重)
- 我们的目标是重现原始论文中的这张表格。
尝试演示笔记本
- 确保您已安装 jupyter notebook。
- 安装 jupyter notebook 的两种替代方式:
# 确保 pip 已升级
pip3 install --upgrade pip
# 安装 jupyter notebook
pip install jupyter
# 在 ssd.pytorch 内运行此命令
jupyter notebook
- 现在导航至
demo/demo.ipynb
,默认地址为 http://localhost:8888,并开始使用吧!
尝试网络摄像头演示
- 支持在 CPU 上运行(可能需要调整
cv2.waitkey
以获得最佳 fps)或在 NVIDIA GPU 上运行 - 此演示当前需要带有 python 绑定的 opencv2+ 和一个板载摄像头
- 您可以在
demo/live.py
中更改默认摄像头
- 您可以在
- 安装 imutils 包以在 CPU 上利用多线程:
pip install imutils
- 运行
python -m demo.live
打开摄像头并开始检测!
TODO
我们已经积累了以下待办事项列表,我们希望在不久的将来完成这些工作:
- 仍待完成:
- 支持 MS COCO 数据集
- 支持 SSD512 训练和测试
- 支持在自定义数据集上进行训练
作者
注意: 不幸的是,这只是我们的业余爱好,而不是全职工作,所以我们会尽力保持更新,但不能保证。也就是说,非常感谢大家的持续帮助和反馈,这对我们来说非常重要。我们会尽快处理所有问题。
参考文献
- Wei Liu 等人。"SSD:单次多框检测器。" ECCV2016。
- 原始实现(CAFFE)
- 特别感谢 Alex Koltun 和他的团队 Webyclip 在完成数据增强部分的帮助。
- 一些其他优秀的 SSD 移植项目,启发来源(特别是 Chainer 仓库):
- [Ch