Este es el repositorio oficial con la implementación en PyTorch de LW-DETR: un reemplazo de transformador para YOLO para detección en tiempo real.
☀️ Si encuentra útil este trabajo para su investigación, por favor destaque nuestro repositorio y cite nuestro artículo. ☀️
Lanzar una serie de modelos de detección en tiempo real en LW-DETR, incluidos LW-DETR-tiny, LW-DETR-small, LW-DETR-medium, LW-DETR-large y LW-DETR-xlarge, denominados
Lanzar una serie de modelos previamente entrenados en LW-DETR. Consulte Hugging Face para descargarlo.
[15/7/2024] Presentamos OVLW-DETR, un eficiente detector de vocabulario abierto con rendimiento sobresaliente y baja latencia, construido sobre LW-DETR. Supera a los detectores de vocabulario abierto en tiempo real existentes en el punto de referencia estándar Zero-Shot LVIS. El código fuente y el modelo previamente entrenado estarán disponibles pronto, ¡estén atentos!
1. Introducción
2. Instalación
3. Preparación
4. entrenar
5. evaluación
6. Implementar
7. Resultados principales
8. Referencias
9. Cita
LW-DETR es un transformador de detección liviano que supera a los YOLO en la detección de objetos en tiempo real. La arquitectura es una pila simple de un codificador ViT, un proyector y un decodificador DETR superficial. LW-DETR aprovecha técnicas avanzadas recientes, como técnicas efectivas de entrenamiento, por ejemplo, pérdida mejorada y preentrenamiento, y ventanas entrelazadas y atenciones globales para reducir la complejidad del codificador ViT. LW-DETR mejora el codificador ViT al agregar mapas de características de varios niveles y los mapas de características intermedias y finales en el codificador ViT, formando mapas de características más ricos e introduce la organización de mapas de características de ventana principal para mejorar la eficiencia del cálculo de atención entrelazada. LW-DETR logra un rendimiento superior que los detectores en tiempo real existentes, por ejemplo, YOLO y sus variantes, en COCO y otros conjuntos de datos de referencia.
El código se desarrolla y valida bajo python=3.8.19, pytorch=1.13.0, cuda=11.6,TensorRT-8.6.1.6 . Es posible que también haya versiones superiores disponibles.
Crea tu propio entorno Python con Anaconda.
conda crear -n lwdetr python = 3.8.19 conda activar lwdetr
Clona este repositorio.
clon de git https://github.com/Atten4Vis/LW-DETR.gitcd LW-DETR
Instale PyTorch y torchvision.
Siga las instrucciones en https://pytorch.org/get-started/locally/.
# un ejemplo: conda install pytorch==1.13.0 torchvision==0.14.0 pytorch-cuda=11.6 -c pytorch -c nvidia
Instale los paquetes necesarios.
Para formación y evaluación:
instalación de pip -r requisitos.txt
Para implementación:
Consulte NVIDIA para obtener instrucciones de instalación de TensorRT.
instalación de pip -r implementar/requisitos.txt
Compilando operadores CUDA
modelos/operaciones de cd python setup.py build install# prueba unitaria (debería ver que todas las comprobaciones son Verdaderas)python test.pycd ../..
Para el conjunto de datos de MS COCO , descargue y extraiga imágenes de trenes y valores de COCO 2017 con anotaciones de http://cocodataset.org. Esperamos que la estructura del directorio sea la siguiente:
COCODIR/ ├── train2017/ ├── val2017/ └── annotations/ ├── instances_train2017.json └── instances_val2017.json
Para el conjunto de datos de Objects365 para entrenamiento previo, descargue imágenes de Objects365 con anotaciones desde https://www.objects365.org/overview.html.
Todos los puntos de control se pueden encontrar en Hugging Face.
Entrenamiento previo en Objects365.
Preentrené el ViT.
Entrenamos previamente el ViT en el conjunto de datos Objects365 utilizando un método MIM, CAE v2, basado en los modelos previamente entrenados. Consulte el siguiente enlace para descargar los modelos previamente entrenados y colocarlos en pretrain_weights/ .
| Modelo | Comentario |
|---|---|
| caev2_tiny_300e_objects365 | modelo ViT previamente entrenado en objetos 365 para LW-DETR-tiny/small usando CAE v2 |
| caev2_tiny_300e_objects365 | modelo ViT previamente entrenado en objetos 365 para LW-DETR-medio/grande usando CAE v2 |
| caev2_tiny_300e_objects365 | modelo ViT previamente entrenado en objetos 365 para LW-DETR-xlarge usando CAE v2 |
LW-DETR previamente entrenado.
Volvemos a entrenar el codificador y entrenamos el proyector y el decodificador en Objects365 de forma supervisada. Consulte el siguiente enlace para descargar los modelos previamente entrenados y colocarlos en pretrain_weights/ .
| Modelo | Comentario |
|---|---|
| LWDETR_tiny_30e_objects365 | modelo LW-DETR-tiny preentrenado en objetos 365 |
| LWDETR_small_30e_objects365 | modelo pequeño LW-DETR previamente entrenado en objetos 365 |
| LWDETR_medium_30e_objects365 | modelo medio LW-DETR previamente entrenado en objetos 365 |
| LWDETR_large_30e_objects365 | modelo grande LW-DETR previamente entrenado en objetos 365 |
| LWDETR_xlarge_30e_objects365 | modelo LW-DETR-xlarge previamente entrenado en objetos 365 |
Ajuste de COCO. Ajustamos el modelo previamente entrenado en COCO. Si desea volver a implementar nuestro repositorio, omita este paso. Si desea evaluar directamente nuestros modelos entrenados, consulte el siguiente enlace para descargar los modelos ajustados y ponerlos en output/ .
| Modelo | Comentario |
|---|---|
| LWDETR_tiny_60e_coco | modelo LW-DETR-tiny ajustado en COCO |
| LWDETR_small_60e_coco | modelo pequeño LW-DETR ajustado en COCO |
| LWDETR_medium_60e_coco | modelo LW-DETR-medio ajustado en COCO |
| LWDETR_large_60e_coco | modelo LW-DETR-grande ajustado en COCO |
| LWDETR_xlarge_60e_coco | modelo LW-DETR-xlarge ajustado en COCO |
Puede ejecutar directamente el archivo scripts/lwdetr_ para el proceso de capacitación en el conjunto de datos coco.
sh scripts/lwdetr_tiny_coco_train.sh /ruta/a/tu/COCODIR
sh scripts/lwdetr_small_coco_train.sh /ruta/a/tu/COCODIR
sh scripts/lwdetr_medium_coco_train.sh /ruta/a/tu/COCODIR
sh scripts/lwdetr_large_coco_train.sh /ruta/a/tu/COCODIR
sh scripts/lwdetr_xlarge_coco_train.sh /ruta/a/tu/COCODIR
Puede ejecutar directamente el archivo scripts/lwdetr_ para el proceso de evaluación en el conjunto de datos coco. Consulte 3. Preparación para descargar una serie de modelos LW-DETR.
sh scripts/lwdetr_tiny_coco_eval.sh /ruta/a/su/COCODIR /ruta/a/su/punto de control
sh scripts/lwdetr_small_coco_eval.sh /ruta/a/su/COCODIR /ruta/a/su/punto de control
sh scripts/lwdetr_medium_coco_eval.sh /ruta/a/su/COCODIR /ruta/a/su/punto de control
sh scripts/lwdetr_large_coco_eval.sh /ruta/a/su/COCODIR /ruta/a/su/punto de control
sh scripts/lwdetr_xlarge_coco_eval.sh /ruta/a/su/COCODIR /ruta/a/su/punto de control
Puede ejecutar el archivo scripts/lwdetr_ para exportar modelos para desarrollo. Antes de la ejecución, asegúrese de que las variables de entorno TensorRT y cuDNN estén configuradas correctamente.
# exportar scripts del modelo ONNX/lwdetr_tiny_coco_export.sh /path/to/your/COCODIR /path/to/your/checkpoint# convertir el modelo de ONNX al motor TensorRT también scripts/lwdetr_tiny_coco_export.sh /path/to/your/COCODIR /path /a/su/punto de control --trt
# exportar scripts del modelo ONNX/lwdetr_small_coco_export.sh /path/to/your/COCODIR /path/to/your/checkpoint# convertir el modelo de ONNX al motor TensorRT también scripts/lwdetr_small_coco_export.sh /path/to/your/COCODIR /path /a/su/punto de control --trt
# exportar scripts del modelo ONNX/lwdetr_medium_coco_export.sh /path/to/your/COCODIR /path/to/your/checkpoint# convertir el modelo de ONNX al motor TensorRT también scripts/lwdetr_medium_coco_export.sh /path/to/your/COCODIR /path /a/su/punto de control --trt
# exportar scripts del modelo ONNX/lwdetr_large_coco_export.sh /path/to/your/COCODIR /path/to/your/checkpoint# convertir el modelo de ONNX al motor TensorRT también scripts/lwdetr_large_coco_export.sh /path/to/your/COCODIR /path /a/su/punto de control --trt
# exportar scripts del modelo ONNX/lwdetr_xlarge_coco_export.sh /path/to/your/COCODIR /path/to/your/checkpoint# convertir el modelo de ONNX al motor TensorRT también scripts/lwdetr_xlarge_coco_export.sh /path/to/your/COCODIR /path /a/su/punto de control --trt
Puede utilizar la herramienta deploy/benchmark.py para ejecutar pruebas comparativas de latencia de inferencia.
# evaluar y comparar la latencia en un modelo onnxpython implementar/benchmark.py --path=/path/to/your/onnxmodel --coco_path=/path/to/your/COCODIR --run_benchmark # evaluar y comparar la latencia en un motor TensorRTpython implementar/benchmark.py --path=/path/to/your/trtengine --coco_path=/path/to/your/COCODIR --run_benchmark
Los principales resultados del conjunto de datos de coco. Informamos el mAP tal como se informó en el artículo original, así como el mAP obtenido de la reimplementación.
Método | preentrenamiento | Parámetros (M) | FLOP (G) | Latencia del modelo (ms) | Latencia total (ms) | mapa | Descargar |
|---|---|---|---|---|---|---|---|
LW-DETR-tiny | ✔ | 12.1 | 11.2 | 2.0 | 2.0 | 42,6 (42,9) | Enlace |
LW-DETR-small | ✔ | 14.6 | 16.6 | 2.9 | 2.9 | 48,0 (48,1) | Enlace |
LW-DETR-medium | ✔ | 28.2 | 42,8 | 5.6 | 5.6 | 52,5 (52,6) | Enlace |
LW-DETR-large | ✔ | 46,8 | 71,6 | 8.8 | 8.8 | 56,1 (56,1) | Enlace |
LW-DETR-xlarge | ✔ | 118.0 | 174,2 | 19.1 | 19.1 | 58,3 (58,3) | Enlace |
Nuestro proyecto se lleva a cabo en base al siguiente documento público con código:
Grupo DETR
ViTDet
DETR deformable
DETR condicional
YOLOv8
AlinearDETR
CAE v2
CAE
Si encuentra útil este código en su investigación, considere citar nuestro artículo:
@article{chen2024lw,title={LW-DETR: Un reemplazo de transformador para YOLO para detección en tiempo real},author={Chen, Qiang y Su, Xiangbo y Zhang, Xinyu y Wang, Jian y Chen, Jiahui y Shen, Yunpeng y Han, Chuchu y Chen, Ziliang y Xu, Weixiang y Li, Fanrong y otros},journal={arXiv preprint arXiv:2406.03459},año={2024}} 
