使用PyTorch实现目标检测新范式DETR(基于transformer)| 留言送书
磐创AI分享
磐创AI分享
【导读】 DETR ( DE tection TR ansformer)的PyTorch训练代码和预训练模型。 我们用Transformer 代替了整个复杂的手工物体检测管道,并用ResNet-50匹配了Faster R-CNN,使用一半的计算能力(FLOP)和相同数量的参数在COCO上获得了 42 AP 。 在PyTorch的50行中进行推断。

这是什么
与传统的计算机视觉技术不同,DETR将目标检测作为一个直接的集合预测问题来处理。它由一个基于集合的全局损失和一个Transformer encoder-decoder 结构组成,该全局损失通过二分匹配强制进行唯一预测。给定固定的学习对象查询集,则DETR会考虑对象与全局图像上下文之间的关系,以直接并行并行输出最终的预测集。由于这种并行性,DETR非常快速和高效。
关于代码
我们认为,物体检测不应该比分类更困难,并且不应要求使用复杂的库来进行训练和推理。DETR的实现和试验非常简单,我们提供了一个[独立Colab笔记本]( https://colab.research.google.com/github/facebook research/detr/blob/colab/notebooks/detr_demo.ipynb )演示如何使用DETR进行推断,这仅需几行PyTorch代码。训练代码遵循这个想法–它不是一个库,而是一个带有标准训练循环的main.py导入模型和条件定义。
有关详细信息,请参见Nicolas Carion, Francisco Massa,Gabriel Synnaeve,Nicolas Usunier,Alexander Kirillov, 和Sergey Zagoruyko.的End-to-End Object Detection with Transformers( https://ai.facebook.com/research/publications/end-to-end-object-detection-with-transformers )
我们提供基准 DETR 和 DETR-DC5 模型,并计划在将来包括更多模型。AP 是使用Torcscript transformer,在COCO 2017 val5k上计算得出的,并且推理时间超过了前100幅val5k COCO图像。
name | backbone | schedule | inf_time | box AP | url | size | |
---|---|---|---|---|---|---|---|
0 | DETR | R50 | 500 | 0.036 | 42.0 | download | 159Mb |
1 | DETR-DC5 | R50 | 500 | 0.083 | 43.3 | download | 159Mb |
2 | DETR | R101 | 500 | 0.050 | 43.5 | download | 232Mb |
3 | DETR-DC5 | R101 | 500 | 0.097 | 44.9 | download | 232Mb |
COCO val5k数据集的评估结果可以在这里gist( https://gist.github.com/szagoruyko/9c9ebb8455610958f7deaa27845d7918 )找到.
COCO全景val5k模型:
name | backbone | box AP | segm AP | PQ | url | size | |
---|---|---|---|---|---|---|---|
0 | DETR | R50 | 38.8 | 31.1 | 43.4 | download | 165Mb |
1 | DETR-DC5 | R50 | 40.2 | 31.9 | 44.6 | download | 165Mb |
2 | DETR | R101 | 40.1 | 33 | 45.1 | download | 237Mb |
这些模型也可通过torch hub找到,以用预训练的权重加载DETR R50,只需执行以下操作:
model = torch.hub.load('facebookresearch/detr', 'detr_resnet50', pretrained=True)
用法
DETR中没有多余的已编译组件,并且程序包依赖项很小,因此代码使用非常简单。我们提供如何通过conda安装依赖项的说明,因此该代码非常易于使用。我们提供了有关如何通过conda安装依赖项的说明。首先,在本地克隆存储库:
git clone https://github.com/facebookresearch/detr.git
然后安装 PyTorch 1.5+ 和 torchvision 0.6+:
conda install -c pytorch pytorch torchvision
安装 pycocotools (用于 COCO 的评估) 和 scipy (用于训练):
conda install cython scipy pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'
这有助于训练和评估检测模型。
(可选)与panoptic一起使用,安装panopticapi:
pip install git+https://github.com/cocodataset/panopticapi.git
数据准备
从 http://cocodataset.org 下载并提取带有标注的COCO 2017训练和验证图像,我们希望目录结构如下:
path/to/coco/ annotations/ # 标注json文件 train2017/ # 训练集图片 val2017/ # 验证集图片
训练
在8个GPU的单个节点上训练基准 DETR 300个 epoch,运行:
python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --coco_path /path/to/coco
一个epoch需要28分钟,因此在一台装有8个V100卡的机器上,300个epoch的训练大约需要6天。为了简化对我们实验结果的复现,我们提供了150个epoch schedule(在一台机器上训练3天)的 结果和训练日志( https://gist.github.com/szagoruyko/b4c3b2c3627294fc369b899987385a3f ),结果达到了39.5 / 60.3 AP / AP50。
我们使用AdamW训练DETR,将transformer中的学习率设置为backbone中的1e-4和1e-5。使用水平翻转,缩放图片来进行图片增强。图像被重新缩放为具有最小800和最大1333的大小。对transformer进行了dropout为0.1的训练,对整个模型进行了0.1的梯度裁剪。
评估
在COCO val5k数据集上使用单个GPU来评估DETR R50:
python main.py --batch_size 2 --no_aux_loss --eval --resume https://dl.fbaipublicfiles.com/detr/detr-r50-e632da11.pth --coco_path /path/to/coco
在gist上,我们提供所有DETR检测模型的结果 。
请注意,数量根据每个GPU的批处理大小(batch size)(图像数量)而有所不同。非DC5模型的批处理大小为2,而DC5的批处理大小为1,因此,如果每个GPU不止对1个图像进行评估,那么DC5模型的AP会显着下降。
模型训练
可以通过Slurm和submitit( https://github.com/facebookincubator/submitit )来进行分布式训练:
pip install submitit
在4个节点上训练基本模型DETR-6-6,300个epoch:
python run_with_submitit.py --timeout 3000 --coco_path /path/to/coco
许可证
DETR是在Apache2.0许可下发布的。请查看LICENSE( https://github.com/facebookresearch/detr/blob/master/LICENSE )文件来获得更多信息。
帮助
我们非常欢迎您的请求!请查看:
CONTRIBUTING.md
( https://github.com/facebookresearch/detr/blob/master/.github/CONTRIBUTING.md )
CODE_OF_CONDUCT.md
( https://github.com/facebookresearch/detr/blob/master/.github/CODE_OF_CONDUCT.md )
留言送书福利
为了感谢大家长期以来的支持,磐创AI推出了“留言送书”活动~在本文文末 留言 即可参与活动,留言内容需为主题相关 。欢迎大家在日常推文中留言,以后将不定期推出 “留言送书”活动。
这次磐小仙精心挑选了一本《 Python数据分析与可视化从入门到精通 》送给 1位粉丝 。
/ 今日留言主题/
谈谈PyTorch的优缺点?
《Python数据分析与可视化从入门到精通》以“零基础”为起点,系统地介绍了Python在数据处理与可视化分析方面的应用,适合从事Python数据处理与可视化的用户学习。
未中奖的朋友也可上京东购买,满100减50活动等你哦!
(点击【阅读原文】查看书籍详情)
本次“留言送书”活动将在6月4号(具体时间以当天推文发出为准,通常为00:15)抽 留言最走心的一位粉丝 免费赠送 这本书籍(走心留言将经过筛选)。届时小仙会公布中奖者评论截图及福利领取方式~
– End –
✄————————————————
看到这里,说明你喜欢这篇文章,请点击「 在看 」或顺手「 转发 」「 点赞 」。
欢迎微信搜索「 panchuangxx 」,添加小编 磐小小仙 微信,每日朋友圈更新一篇高质量推文(无广告),为您提供更多精彩内容。
▼ ▼ 扫描二维码添加小编 ▼ ▼