这是 LW-DETR 的 PyTorch 实现的官方存储库:用于实时检测的 YOLO 的变压器替代品。
☀️ 如果您发现这项工作对您的研究有用,请为我们的存储库加注星标并引用我们的论文! ☀️
在LW-DETR中发布一系列实时检测模型,包括LW-DETR-tiny、LW-DETR-small、LW-DETR-medium、LW-DETR-large和LW-DETR-xlarge,命名为<LWDETR_*size_60e_coco .pth>。请参考抱脸下载。
在LW-DETR中发布一系列预训练模型。请参考抱脸下载。
[2024/7/15] 我们推出了 OVLW-DETR,这是一种基于 LW-DETR 构建的高效开放词汇检测器,具有出色的性能和低延迟。它在标准零样本 LVIS 基准上超越了现有的实时开放词汇检测器。源代码和预训练模型即将推出,敬请期待!
一、简介
2. 安装
3. 准备工作
4. 火车
5. 评估
6. 部署
七、主要结果
8. 参考文献
9. 引文
LW-DETR 是一种轻量级检测变压器,其实时目标检测性能优于 YOLO。该架构是 ViT 编码器、投影仪和浅层 DETR 解码器的简单堆栈。 LW-DETR 利用最新的先进技术,例如训练有效的技术,例如改进的损失和预训练,以及用于降低 ViT 编码器复杂性的交错窗口和全局注意力。 LW-DETR通过聚合多级特征图以及ViT编码器中的中间和最终特征图来改进ViT编码器,形成更丰富的特征图,并引入窗口主特征图组织以提高交错注意力计算的效率。 LW-DETR 在 COCO 和其他基准数据集上实现了比现有实时检测器(例如 YOLO 及其变体)更优越的性能。
该代码是在python=3.8.19, pytorch=1.13.0, cuda=11.6,TensorRT-8.6.1.6
下开发和验证的。也可能提供更高版本。
使用 Anaconda 创建您自己的 Python 环境。
conda 创建-n lwdetr python=3.8.19 conda 激活 lwdetr
克隆这个仓库。
git 克隆 https://github.com/Atten4Vis/LW-DETR.gitcd LW-DETR
安装 PyTorch 和 torchvision。
按照 https://pytorch.org/get-started/locally/ 上的说明进行操作。
# 示例:conda install pytorch==1.13.0 torchvision==0.14.0 pytorch-cuda=11.6 -c pytorch -c nvidia
安装所需的软件包。
用于培训和评估:
pip install -r 要求.txt
对于部署:
TensorRT的安装说明请参考NVIDIA
pip install -r 部署/requirements.txt
编译 CUDA 运算符
CD 型号/操作 python setup.py build install# 单元测试(应该看到所有检查都是 True)python test.pycd ../..
对于MS COCO数据集,请从 http://cocodataset.org 下载并提取带有注释的 COCO 2017 训练和验证图像。我们期望目录结构如下:
COCODIR/ ├── train2017/ ├── val2017/ └── annotations/ ├── instances_train2017.json └── instances_val2017.json
对于用于预训练的Objects365数据集,请从 https://www.objects365.org/overview.html 下载带有注释的 Objects365 图像。
所有关卡都可以在抱脸中找到。
在 Objects365 上进行预训练。
预训练 ViT。
我们基于预训练模型,使用 MIM 方法 CAE v2 在数据集 Objects365 上预训练 ViT。请参考以下链接下载预训练模型,并将其放入pretrain_weights/
中。
模型 | 评论 |
---|---|
caev2_tiny_300e_objects365 | 使用 CAE v2 在objects365 上为 LW-DETR-tiny/small 预训练 ViT 模型 |
caev2_tiny_300e_objects365 | 使用 CAE v2 在objects365 上针对 LW-DETR-中/大型预训练 ViT 模型 |
caev2_tiny_300e_objects365 | 使用 CAE v2 在objects365 上为 LW-DETR-xlarge 预训练 ViT 模型 |
预训练的 LW-DETR。
我们以监督方式在 Objects365 上重新训练编码器并训练投影仪和解码器。请参考以下链接下载预训练模型,并将其放入pretrain_weights/
中。
模型 | 评论 |
---|---|
LWDETR_tiny_30e_objects365 | 在对象上预训练的 LW-DETR-tiny 模型365 |
LWDETR_small_30e_objects365 | 在物体上预训练的 LW-DETR-小模型365 |
LWDETR_medium_30e_objects365 | 在对象上预训练的 LW-DETR-medium 模型365 |
LWDETR_large_30e_objects365 | 在对象上预训练的 LW-DETR-large 模型365 |
LWDETR_xlarge_30e_objects365 | 在对象上预训练的 LW-DETR-xlarge 模型365 |
对 COCO 进行微调。我们对 COCO 上的预训练模型进行了微调。如果您想重新实现我们的存储库,请跳过此步骤。如果您想直接评估我们训练好的模型,请参考以下链接下载微调后的模型,并将其放入output/
中。
模型 | 评论 |
---|---|
LWDETR_tiny_60e_coco | COCO 上微调的 LW-DETR-tiny 模型 |
LWDETR_small_60e_coco | COCO 上微调的 LW-DETR-小模型 |
LWDETR_medium_60e_coco | COCO 上微调的 LW-DETR-medium 模型 |
LWDETR_large_60e_coco | COCO 上微调的 LW-DETR-large 模型 |
LWDETR_xlarge_60e_coco | COCO 上微调的 LW-DETR-xlarge 模型 |
您可以直接运行scripts/lwdetr_<model_size>_coco_train.sh
文件在coco数据集上进行训练过程。
sh 脚本/lwdetr_tiny_coco_train.sh /path/to/your/COCODIR
sh 脚本/lwdetr_small_coco_train.sh /path/to/your/COCODIR
sh 脚本/lwdetr_medium_coco_train.sh /path/to/your/COCODIR
sh 脚本/lwdetr_large_coco_train.sh /path/to/your/COCODIR
sh 脚本/lwdetr_xlarge_coco_train.sh /path/to/your/COCODIR
您可以直接运行scripts/lwdetr_<model_size>_coco_eval.sh
文件来对coco数据集进行评估过程。请参考3.准备下载LW-DETR系列模型。
sh 脚本/lwdetr_tiny_coco_eval.sh /path/to/your/COCODIR /path/to/your/checkpoint
sh 脚本/lwdetr_small_coco_eval.sh /path/to/your/COCODIR /path/to/your/checkpoint
sh 脚本/lwdetr_medium_coco_eval.sh /path/to/your/COCODIR /path/to/your/checkpoint
sh 脚本/lwdetr_large_coco_eval.sh /path/to/your/COCODIR /path/to/your/checkpoint
sh 脚本/lwdetr_xlarge_coco_eval.sh /path/to/your/COCODIR /path/to/your/checkpoint
您可以运行scripts/lwdetr_<model_size>_coco_export.sh
文件来导出模型以进行开发。执行之前,请确保TensorRT和cuDNN环境变量设置正确。
# 导出 ONNX modelsh script/lwdetr_tiny_coco_export.sh /path/to/your/COCODIR /path/to/your/checkpoint# 将模型从 ONNX 转换为 TensorRT 引擎 wellsh script/lwdetr_tiny_coco_export.sh /path/to/your/COCODIR /path /到/你的/检查点--trt
# 导出 ONNX modelsh script/lwdetr_small_coco_export.sh /path/to/your/COCODIR /path/to/your/checkpoint# 将模型从 ONNX 转换为 TensorRT 引擎 wellsh script/lwdetr_small_coco_export.sh /path/to/your/COCODIR /path /到/你的/检查点--trt
# 导出 ONNX modelsh 脚本/lwdetr_medium_coco_export.sh /path/to/your/COCODIR /path/to/your/checkpoint# 将模型从 ONNX 转换为 TensorRT 引擎 wellsh 脚本/lwdetr_medium_coco_export.sh /path/to/your/COCODIR /path /到/你的/检查点--trt
# 导出 ONNX modelsh script/lwdetr_large_coco_export.sh /path/to/your/COCODIR /path/to/your/checkpoint# 将模型从 ONNX 转换为 TensorRT 引擎以及sh script/lwdetr_large_coco_export.sh /path/to/your/COCODIR /path /到/你的/检查点--trt
# 导出 ONNX modelsh script/lwdetr_xlarge_coco_export.sh /path/to/your/COCODIR /path/to/your/checkpoint# 将模型从 ONNX 转换为 TensorRT 引擎 wellsh script/lwdetr_xlarge_coco_export.sh /path/to/your/COCODIR /path /到/你的/检查点--trt
您可以使用deploy/benchmark.py
工具来运行推理延迟的基准测试。
# 评估 onnx 模型上的延迟并对其进行基准测试python deploy/benchmark.py --path=/path/to/your/onnxmodel --coco_path=/path/to/your/COCODIR --run_benchmark # 评估 TensorRT 引擎上的延迟并对其进行基准测试python deploy/benchmark.py --path=/path/to/your/trtengine --coco_path=/path/to/your/COCODIR --run_benchmark
coco数据集上的主要结果。我们报告了原始论文中报告的 mAP,以及从重新实现中获得的 mAP。
方法 | 预训练 | 参数 (M) | 失败次数 (G) | 模型延迟(毫秒) | 总延迟(毫秒) | 地图 | 下载 |
---|---|---|---|---|---|---|---|
LW-DETR-tiny | ✔ | 12.1 | 11.2 | 2.0 | 2.0 | 42.6 (42.9) | 关联 |
LW-DETR-small | ✔ | 14.6 | 16.6 | 2.9 | 2.9 | 48.0 (48.1) | 关联 |
LW-DETR-medium | ✔ | 28.2 | 42.8 | 5.6 | 5.6 | 52.5 (52.6) | 关联 |
LW-DETR-large | ✔ | 46.8 | 71.6 | 8.8 | 8.8 | 56.1 (56.1) | 关联 |
LW-DETR-xlarge | ✔ | 118.0 | 174.2 | 19.1 | 19.1 | 58.3 (58.3) | 关联 |
我们的项目是基于以下公开论文和代码进行的:
德特集团
维生素D
可变形DETR
有条件DETR
YOLOv8
对齐DETR
计算机辅助工程 v2
计算机辅助工程
如果您发现此代码对您的研究有用,请考虑引用我们的论文:
@article{chen2024lw,title={LW-DETR: 用于实时检测的 YOLO 的 Transformer 替代品},author={Chen, Qiang and Su, Xianangbo and Zhang, Xinyu and Wang,jian and Chen, Jiahui and Shen, Yunpeng以及韩、楚楚与陈、子良与徐、伟祥与李、繁荣等},journal={arXiv 预印本arXiv:2406.03459},年份={2024}}