detrex: 基于DETR的目标检测研究平台

detrex是一个为DETR(DEtection TRansformer)系列目标检测算法设计的开源研究平台。它由IDEA研究团队开发,旨在为研究人员和开发者提供一个统一、高度模块化、轻量级的代码库,以便于开发和比较不同的DETR变体算法。

主要特性

detrex具有以下几个主要特性:

模块化设计: detrex将基于Transformer的目标检测框架分解为多个组件,使用户可以轻松构建自定义模型。
强大的基线: detrex提供了一系列DETR变体算法的强大基线实现。通过优化超参数,大多数支持的算法性能提升了0.2到1.1个AP。
易于使用: detrex设计为轻量级且易于使用:
- 采用LazyConfig系统,提供更灵活的语法和更清晰的配置文件。
- 基于detectron2修改的轻量级训练引擎。
广泛的算法支持: 支持多种主流的DETR变体算法,包括DETR、Deformable DETR、Conditional DETR、DAB-DETR、DN-DETR、DINO等。
多任务支持: 除了目标检测外,还支持实例分割、姿态估计等视觉识别任务。

安装使用

detrex的安装非常简单,只需要以下几个步骤:

git clone https://github.com/IDEA-Research/detrex.git
cd detrex
pip install -e .

安装完成后,可以通过以下方式使用预训练模型进行推理:

from detrex import build_model
from detrex.config import get_config

# 加载配置文件
cfg = get_config("projects/dino/configs/dino_r50_4scale_12ep.py")

# 构建模型
model = build_model(cfg)

# 加载预训练权重
model.load_state_dict(torch.load("path/to/pretrained_weights.pth"))

# 进行推理
outputs = model(images)

模型训练与评估

detrex提供了方便的命令行工具用于模型训练和评估。以DAB-DETR为例:

# 训练
python tools/train_net.py \
    --config-file projects/dab_detr/configs/dab_detr_r50_50ep.py \
    --num-gpus 8

# 评估
python tools/train_net.py \
    --config-file projects/dab_detr/configs/dab_detr_r50_50ep.py \
    --eval-only train.init_checkpoint=/path/to/checkpoint

支持的算法

detrex目前支持以下DETR变体算法:

DETR
Deformable DETR
Conditional DETR
DAB-DETR
DN-DETR
DINO
Group DETR
H-Deformable-DETR
MaskDINO
Focus-DETR
等等每个算法都有详细的配置文件和预训练模型可供使用。

模型性能

通过优化训练策略和超参数,detrex在多个算法上都取得了不错的性能提升。以下是部分算法在COCO数据集上的检测结果:

算法	主干网络	AP
DINO	Swin-L	58.5
DAB-Deformable-DETR	R50	49.0
H-Deformable-DETR	Swin-L	57.5
MaskDINO	Swin-L	57.0
更多详细的性能指标可以参考Model Zoo。

社区与贡献

detrex是一个开源项目,欢迎社区贡献。您可以通过以下方式参与:

在GitHub Issues中报告bug或提出新功能建议
提交Pull Request来改进代码
在Discussions中讨论使用问题和想法我们期待您的参与,一起推动DETR系列算法的发展!

总结

detrex为DETR系列目标检测算法提供了一个强大而灵活的研究平台。它不仅实现了多种主流算法,还提供了便捷的工具支持实验和分析。无论您是研究人员还是实践者,detrex都是一个值得尝试的优秀工具。希望本文能帮助您了解detrex项目,开始您的DETR算法探索之旅。如果您对项目有任何问题或建议,欢迎在GitHub上与我们交流!