pytorch YOLOv4
1.0.0
Une mise en œuvre minimale de Pytorch de Yolov4.
Paper Yolo V4: https://arxiv.org/abs/2004.10934
Code source: https: //github.com/alexeyab/darknet
Plus de détails: http://pjreddie.com/darknet/yolo/
Inférence
Former
├── README.md
├── dataset.py dataset
├── demo.py demo to run pytorch --> tool/darknet2pytorch
├── demo_darknet2onnx.py tool to convert into onnx --> tool/darknet2pytorch
├── demo_pytorch2onnx.py tool to convert into onnx
├── models.py model for pytorch
├── train.py train models.py
├── cfg.py cfg.py for train
├── cfg cfg --> darknet2pytorch
├── data
├── weight --> darknet2pytorch
├── tool
│ ├── camera.py a demo camera
│ ├── coco_annotation.py coco dataset generator
│ ├── config.py
│ ├── darknet2pytorch.py
│ ├── region_loss.py
│ ├── utils.py
│ └── yolo_layer.py
Vous pouvez utiliser darknet2pytorch pour le convertir vous-même ou télécharger mon modèle converti.
Utilisez Yolov4 pour former vos propres données
Télécharger le poids
Transformer les données
Pour l'ensemble de données CoCo, vous pouvez utiliser Tool / CoCo_annotation.py.
# train.txt
image_path1 x1,y1,x2,y2,id x1,y1,x2,y2,id x1,y1,x2,y2,id ...
image_path2 x1,y1,x2,y2,id x1,y1,x2,y2,id x1,y1,x2,y2,id ...
...
...
Former
Vous pouvez définir des paramètres dans cfg.py.
python train.py -g [GPU_ID] -dir [Dataset direction] ...
Les modèles ONNX et Tensorrt sont convertis à partir de pytorch (tianxiaomo): pytorch-> onnx-> Tensorrt. Voir les sections suivantes pour plus de détails sur les conversions.
| Type de modèle | AP | AP50 | AP75 | SPA | APM | Apogée |
|---|---|---|---|---|---|---|
| Darknet (papier Yolov4) | 0,471 | 0,710 | 0,510 | 0,278 | 0,525 | 0,636 |
| Pytorch (Tianxiaomo) | 0,466 | 0,704 | 0,505 | 0,267 | 0,524 | 0,629 |
| Tensorrt Fp32 + Batchednmsplugin | 0,472 | 0,708 | 0,511 | 0,273 | 0,530 | 0,637 |
| Tensorrt FP16 + Batchednmsplugin | 0,472 | 0,708 | 0,511 | 0,273 | 0,530 | 0,636 |
| Type de modèle | AP | AP50 | AP75 | SPA | APM | Apogée |
|---|---|---|---|---|---|---|
| Darknet (papier Yolov4) | 0,412 | 0,628 | 0,443 | 0,204 | 0,444 | 0,560 |
| Pytorch (Tianxiaomo) | 0,404 | 0,615 | 0,436 | 0.196 | 0,438 | 0,552 |
| Tensorrt Fp32 + Batchednmsplugin | 0,412 | 0,625 | 0,445 | 0,200 | 0,446 | 0,564 |
| Tensorrt FP16 + Batchednmsplugin | 0,412 | 0,625 | 0,445 | 0,200 | 0,446 | 0,563 |
La taille de l'entrée d'image n'est pas limitée en 320 * 320 , 416 * 416 , 512 * 512 et 608 * 608 . Vous pouvez ajuster vos tailles d'entrée pour un rapport d'entrée différent, par exemple: 320 * 608 . Une plus grande taille d'entrée pourrait aider à détecter des cibles plus petites, mais peut être plus lente et la mémoire GPU épuisante.
height = 320 + 96 * n , n in { 0 , 1 , 2 , 3 , ...}
width = 320 + 96 * m , m in { 0 , 1 , 2 , 3 , ...}Chargez le modèle Darknet pré-entraîné et les poids DarkNet pour faire l'inférence (la taille de l'image est déjà configurée dans le fichier CFG)
python demo.py -cfgfile < cfgFile > -weightfile < weightFile > -imgfile < imgFile >Chargez les poids Pytorch (fichier PTH) pour effectuer l'inférence
python models.py < num_classes > < weightfile > < imgfile > < IN_IMAGE_H > < IN_IMAGE_W > < namefile(optional) >Chargez le fichier ONNX converti pour faire l'inférence (voir section 3 et 4)
Chargez le fichier moteur Tensorrt converti pour effectuer une inférence (voir section 5)
Il y a 2 sorties d'inférence.
[batch, num_boxes, 1, 4] qui représente x1, y1, x2, y2 de chaque boîte de délimitation.[batch, num_boxes, num_classes] indiquant des scores de toutes les classes pour chaque boîte de délimitation.Jusqu'à présent, un petit morceau de post-traitement, y compris le NMS, est nécessaire. Nous essayons de minimiser le temps et la complexité du post-traitement.
Ce script doit convertir le modèle Darknet officiel prétrainé en onnx
Version pytorch recommandée:
Installer onnxruntime
pip install onnxruntimeExécutez le script Python pour générer un modèle ONNX et exécuter la démo
python demo_darknet2onnx.py < cfgFile > < namesFile > < weightFile > < imageFile > < batchSize > Vous pouvez convertir votre modèle Pytorch formé en onnx en utilisant ce script
Version pytorch recommandée:
Installer onnxruntime
pip install onnxruntimeExécutez le script Python pour générer un modèle ONNX et exécuter la démo
python demo_pytorch2onnx.py < weight_file > < image_path > < batch_size > < n_classes > < IN_IMAGE_H > < IN_IMAGE_W >Par exemple:
python demo_pytorch2onnx.py yolov4.pth dog.jpg 8 80 416 416Exécutez la commande suivante pour convertir le modèle Yolov4 ONNX en moteur Tensorrt
trtexec --onnx= < onnx_file > --explicitBatch --saveEngine= < tensorRT_engine_file > --workspace= < size_in_megabytes > --fp16Exécutez la commande suivante pour convertir le modèle Yolov4 ONNX en moteur Tensorrt
trtexec --onnx= < onnx_file >
--minShapes=input: < shape_of_min_batch > --optShapes=input: < shape_of_opt_batch > --maxShapes=input: < shape_of_max_batch >
--workspace= < size_in_megabytes > --saveEngine= < engine_file > --fp16Par exemple:
trtexec --onnx=yolov4_-1_3_320_512_dynamic.onnx
--minShapes=input:1x3x320x512 --optShapes=input:4x3x320x512 --maxShapes=input:8x3x320x512
--workspace=2048 --saveEngine=yolov4_-1_3_320_512_dynamic.engine --fp16python demo_trt.py < tensorRT_engine_file > < input_image > < input_H > < input_W >Cette démo ne fonctionne ici que lorsque BatchSize est dynamique (1 devrait être dans la plage dynamique) ou BatchSize = 1, mais vous pouvez mettre à jour cette démo un peu pour d'autres tailles de lots dynamiques ou statiques.
Remarque1: Input_h et Input_w doivent être d'accord avec la taille d'entrée dans le fichier ONNX d'origine.
Remarque2: Des opérations NMS supplémentaires sont nécessaires pour la sortie Tensorrt. Cette démo utilise le code NMS Python à partir de tool/utils.py .
Premièrement: conversion en onnx
TensorFlow> = 2.0
1: Merci: github: https: //github.com/onnx/onnx-tensorflow
2: Exécutez Git Clone https://github.com/onnx/onnx-tetensorflow.git && cd onnx-tensorflow run pip install -e.
Remarque: les erreurs se produiront lors de l'utilisation de "PIP d'installation ONNX-TF", du moins pour moi, il est recommandé d'utiliser l'installation du code source
cd DeepStream
make
Pour un seul lot,
trtexec --onnx=<onnx_file> --explicitBatch --saveEngine=<tensorRT_engine_file> --workspace=<size_in_megabytes> --fp16
Pour multi-lots,
trtexec --onnx=<onnx_file> --explicitBatch --shapes=input:Xx3xHxW --optShapes=input:Xx3xHxW --maxShapes=input:Xx3xHxW --minShape=input:1x3xHxW --saveEngine=<tensorRT_engine_file> --fp16
Remarque: les maux de maxais ne peuvent pas être plus grands que le modèle de forme d'origine.
Référence:
@article{yolov4,
title={YOLOv4: YOLOv4: Optimal Speed and Accuracy of Object Detection},
author={Alexey Bochkovskiy, Chien-Yao Wang, Hong-Yuan Mark Liao},
journal = {arXiv},
year={2020}
}
