Este é o repositório oficial com implementação PyTorch de LW-DETR: uma substituição de transformador para YOLO para detecção em tempo real.
☀️ Se você achar este trabalho útil para sua pesquisa, por favor, marque nosso repositório com estrela e cite nosso artigo! ☀️
Lance uma série de modelos de detecção em tempo real em LW-DETR, incluindo LW-DETR-tiny, LW-DETR-small, LW-DETR-medium, LW-DETR-large e LW-DETR-xlarge, denominados
Lance uma série de modelos pré-treinados em LW-DETR. Consulte Hugging Face para fazer o download.
[2024/7/15] Apresentamos o OVLW-DETR, um detector de vocabulário aberto eficiente com excelente desempenho e baixa latência, baseado em LW-DETR. Ele supera os detectores de vocabulário aberto em tempo real existentes no benchmark Zero-Shot LVIS padrão. O código-fonte e o modelo pré-treinado serão lançados em breve, fique ligado!
1. Introdução
2. Instalação
3. Preparação
4. Treine
5. Avaliação
6. Implantar
7. Principais Resultados
8. Referências
9. Citação
LW-DETR é um transformador de detecção leve, que supera os YOLOs para detecção de objetos em tempo real. A arquitetura é uma pilha simples de um codificador ViT, um projetor e um decodificador DETR raso. O LW-DETR aproveita técnicas avançadas recentes, como técnicas eficazes de treinamento, por exemplo, perda e pré-treinamento aprimorados, e janela intercalada e atenções globais para reduzir a complexidade do codificador ViT. LW-DETR melhora o codificador ViT agregando mapas de recursos de vários níveis e os mapas de recursos intermediários e finais no codificador ViT, formando mapas de recursos mais ricos, e introduz a organização do mapa de recursos de janela principal para melhorar a eficiência da computação de atenção intercalada. O LW-DETR alcança desempenho superior ao dos detectores em tempo real existentes, por exemplo, YOLO e suas variantes, no COCO e outros conjuntos de dados de referência.
O código é desenvolvido e validado em python=3.8.19, pytorch=1.13.0, cuda=11.6,TensorRT-8.6.1.6 . Versões superiores também podem estar disponíveis.
Crie seu próprio ambiente Python com Anaconda.
conda criar -n lwdetr python=3.8.19 conda ativar lwdetr
Clone este repositório.
clone git https://github.com/Atten4Vis/LW-DETR.gitcd LW-DETR
Instale PyTorch e torchvision.
Siga as instruções em https://pytorch.org/get-started/locally/.
# um exemplo:conda install pytorch==1.13.0 torchvision==0.14.0 pytorch-cuda=11.6 -c pytorch -c nvidia
Instale os pacotes necessários.
Para treinamento e avaliação:
pip instalar -r requisitos.txt
Para implantação:
Consulte a NVIDIA para obter instruções de instalação do TensorRT
pip install -r implantar/requisitos.txt
Compilando operadores CUDA
modelos/operações de CD python setup.py build install# unit test (deve ver que toda a verificação é True)python test.pycd ../..
Para o conjunto de dados MS COCO , baixe e extraia imagens de trem e val COCO 2017 com anotações de http://cocodataset.org. Esperamos que a estrutura de diretórios seja a seguinte:
COCODIR/ ├── train2017/ ├── val2017/ └── annotations/ ├── instances_train2017.json └── instances_val2017.json
Para o conjunto de dados Objects365 para pré-treinamento, baixe imagens Objects365 com anotações em https://www.objects365.org/overview.html.
Todos os pontos de verificação podem ser encontrados em Hugging Face.
Pré-treinamento em Objects365.
Pré-treinei o ViT.
Pré-treinamos o ViT no conjunto de dados Objects365 usando um método MIM, CAE v2, baseado nos modelos pré-treinados. Consulte o link a seguir para baixar os modelos pré-treinados e colocá-los em pretrain_weights/ .
| Modelo | Comentário |
|---|---|
| caev2_tiny_300e_objects365 | modelo ViT pré-treinado em objects365 para LW-DETR-tiny/small usando CAE v2 |
| caev2_tiny_300e_objects365 | modelo ViT pré-treinado em objects365 para LW-DETR-médio/grande usando CAE v2 |
| caev2_tiny_300e_objects365 | modelo ViT pré-treinado em objects365 para LW-DETR-xlarge usando CAE v2 |
LW-DETR pré-treinado.
Treinamos novamente o codificador e treinamos o projetor e o decodificador no Objects365 de forma supervisionada. Consulte o link a seguir para baixar os modelos pré-treinados e colocá-los em pretrain_weights/ .
| Modelo | Comentário |
|---|---|
| LWDETR_tiny_30e_objects365 | modelo LW-DETR-minúsculo pré-treinado em objetos365 |
| LWDETR_small_30e_objects365 | modelo LW-DETR-pequeno pré-treinado em objetos365 |
| LWDETR_medium_30e_objects365 | modelo pré-treinado de meio LW-DETR em objetos365 |
| LWDETR_large_30e_objects365 | modelo LW-DETR-grande pré-treinado em objetos365 |
| LWDETR_xlarge_30e_objects365 | modelo LW-DETR-xlarge pré-treinado em objetos365 |
Ajuste fino no COCO. Ajustamos o modelo pré-treinado no COCO. Se você deseja reimplementar nosso repositório, pule esta etapa. Se você deseja avaliar diretamente nossos modelos treinados, consulte o link a seguir para baixar os modelos ajustados e colocá-los em output/ .
| Modelo | Comentário |
|---|---|
| LWDETR_tiny_60e_coco | modelo LW-DETR-miny ajustado no COCO |
| LWDETR_small_60e_coco | modelo LW-DETR-pequeno ajustado no COCO |
| LWDETR_medium_60e_coco | modelo LW-DETR médio ajustado no COCO |
| LWDETR_large_60e_coco | modelo LW-DETR-grande ajustado no COCO |
| LWDETR_xlarge_60e_coco | modelo LW-DETR-xlarge ajustado no COCO |
Você pode executar diretamente o arquivo scripts/lwdetr_ para o processo de treinamento no conjunto de dados coco.
sh scripts/lwdetr_tiny_coco_train.sh /caminho/para/seu/COCODIR
sh scripts/lwdetr_small_coco_train.sh /caminho/para/seu/COCODIR
sh scripts/lwdetr_medium_coco_train.sh /caminho/para/seu/COCODIR
sh scripts/lwdetr_large_coco_train.sh /caminho/para/seu/COCODIR
sh scripts/lwdetr_xlarge_coco_train.sh /caminho/para/seu/COCODIR
Você pode executar diretamente o arquivo scripts/lwdetr_ para o processo de avaliação no conjunto de dados coco. Consulte 3. Preparação para baixar uma série de modelos LW-DETR.
sh scripts/lwdetr_tiny_coco_eval.sh /caminho/para/seu/COCODIR /caminho/para/seu/ponto de verificação
sh scripts/lwdetr_small_coco_eval.sh /caminho/para/seu/COCODIR /caminho/para/seu/ponto de verificação
sh scripts/lwdetr_medium_coco_eval.sh /caminho/para/seu/COCODIR /caminho/para/seu/ponto de verificação
sh scripts/lwdetr_large_coco_eval.sh /caminho/para/seu/COCODIR /caminho/para/seu/ponto de verificação
sh scripts/lwdetr_xlarge_coco_eval.sh /caminho/para/seu/COCODIR /caminho/para/seu/ponto de verificação
Você pode executar o arquivo scripts/lwdetr_ para exportar modelos para desenvolvimento. Antes da execução, certifique-se de que as variáveis de ambiente TensorRT e cuDNN estejam definidas corretamente.
# exportar scripts de modelos ONNX/lwdetr_tiny_coco_export.sh /path/to/your/COCODIR /path/to/your/checkpoint# converter modelo de ONNX para motor TensorRT como scripts wellsh/lwdetr_tiny_coco_export.sh /path/to/your/COCODIR /path /para/seu/ponto de verificação --trt
# exportar scripts de modelos ONNX/lwdetr_small_coco_export.sh /path/to/your/COCODIR /path/to/your/checkpoint# converter modelo de ONNX para mecanismo TensorRT como scripts wellsh/lwdetr_small_coco_export.sh /path/to/your/COCODIR /path /para/seu/ponto de verificação --trt
# exportar scripts de modelos ONNX/lwdetr_medium_coco_export.sh /path/to/your/COCODIR /path/to/your/checkpoint# converter modelo de ONNX para motor TensorRT como scripts wellsh/lwdetr_medium_coco_export.sh /path/to/your/COCODIR /path /para/seu/ponto de verificação --trt
# exportar scripts de modelos ONNX/lwdetr_large_coco_export.sh /path/to/your/COCODIR /path/to/your/checkpoint# converter modelo de ONNX para motor TensorRT como scripts wellsh/lwdetr_large_coco_export.sh /path/to/your/COCODIR /path /para/seu/ponto de verificação --trt
# exportar scripts de modelos ONNX/lwdetr_xlarge_coco_export.sh /path/to/your/COCODIR /path/to/your/checkpoint# converter modelo de ONNX para motor TensorRT como scripts wellsh/lwdetr_xlarge_coco_export.sh /path/to/your/COCODIR /path /para/seu/ponto de verificação --trt
Você pode usar a ferramenta deploy/benchmark.py para executar benchmarks de latência de inferência.
# avalia e compara a latência em um modelo onnxpython deploy/benchmark.py --path=/path/to/your/onnxmodel --coco_path=/path/to/your/COCODIR --run_benchmark # avaliar e comparar a latência em um mecanismo TensorRTpython deploy/benchmark.py --path=/path/to/your/trtengine --coco_path=/path/to/your/COCODIR --run_benchmark
Os principais resultados no conjunto de dados coco. Relatamos o mAP conforme relatado no artigo original, bem como o mAP obtido na reimplementação.
Método | pré-treinamento | Parâmetros (M) | FLOPs (G) | Latência do modelo (ms) | Latência total (ms) | mapa | Download |
|---|---|---|---|---|---|---|---|
LW-DETR-tiny | ✔ | 12.1 | 11.2 | 2,0 | 2,0 | 42,6 (42,9) | Link |
LW-DETR-small | ✔ | 14.6 | 16.6 | 2.9 | 2.9 | 48,0 (48,1) | Link |
LW-DETR-medium | ✔ | 28.2 | 42,8 | 5.6 | 5.6 | 52,5 (52,6) | Link |
LW-DETR-large | ✔ | 46,8 | 71,6 | 8.8 | 8.8 | 56,1 (56,1) | Link |
LW-DETR-xlarge | ✔ | 118,0 | 174,2 | 19.1 | 19.1 | 58,3 (58,3) | Link |
Nosso projeto é conduzido com base no seguinte documento público com código:
Grupo DETR
ViTDet
DETR deformável
DETR condicional
YOLOv8
AlinharDETR
CAE v2
CAE
Se você achar este código útil em sua pesquisa, considere citar nosso artigo:
@artigo{chen2024lw,title={LW-DETR: Uma substituição de transformador para YOLO para detecção em tempo real},autor={Chen, Qiang e Su, Xiangbo e Zhang, Xinyu e Wang, Jian e Chen, Jiahui e Shen, Yunpeng e Han, Chuchu e Chen, Ziliang e Xu, Weixiang e Li, Fanrong e outros},journal={arXiv preprint arXiv:2406.03459},ano={2024}} 
