pytorch YOLOv4
1.0.0
Implementasi Pytorch minimal YOLOV4.
Paper Yolo V4: https://arxiv.org/abs/2004.10934
Kode Sumber: https: //github.com/alexeyab/darknet
Rincian lebih lanjut: http://pjreddie.com/darknet/yolo/
Kesimpulan
Kereta
├── 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
Anda dapat menggunakan DarkNet2PyTorch untuk mengonversinya sendiri, atau mengunduh model saya yang dikonversi.
Gunakan YOLOV4 untuk melatih data Anda sendiri
Unduh Berat Badan
Mengubah data
Untuk dataset Coco, Anda dapat menggunakan alat/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 ...
...
...
Kereta
Anda dapat mengatur parameter di cfg.py.
python train.py -g [GPU_ID] -dir [Dataset direction] ...
Model ONNX dan TensorRT dikonversi dari Pytorch (Tianxiaomo): Pytorch-> onnx-> Tensorrt. Lihat bagian berikut untuk detail konversi lebih lanjut.
| Tipe model | Ap | AP50 | Ap75 | APS | APM | Apl |
|---|---|---|---|---|---|---|
| Darknet (kertas 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 |
| Tipe model | Ap | AP50 | Ap75 | APS | APM | Apl |
|---|---|---|---|---|---|---|
| Darknet (kertas 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 |
Ukuran input gambar tidak dibatasi dalam 320 * 320 , 416 * 416 , 512 * 512 dan 608 * 608 . Anda dapat menyesuaikan ukuran input Anda untuk rasio input yang berbeda, misalnya: 320 * 608 . Ukuran input yang lebih besar dapat membantu mendeteksi target yang lebih kecil, tetapi mungkin lebih lambat dan memori GPU melelahkan.
height = 320 + 96 * n , n in { 0 , 1 , 2 , 3 , ...}
width = 320 + 96 * m , m in { 0 , 1 , 2 , 3 , ...}Muat model darknet pretrained dan bobot darknet untuk melakukan inferensi (ukuran gambar sudah dikonfigurasi dalam file cfg)
python demo.py -cfgfile < cfgFile > -weightfile < weightFile > -imgfile < imgFile >Muat bobot pytorch (file PTH) untuk melakukan inferensi
python models.py < num_classes > < weightfile > < imgfile > < IN_IMAGE_H > < IN_IMAGE_W > < namefile(optional) >Load File Onnx yang dikonversi untuk melakukan inferensi (lihat Bagian 3 dan 4)
Load File Mesin Tensorrt yang dikonversi untuk melakukan inferensi (lihat Bagian 5)
Ada 2 output inferensi.
[batch, num_boxes, 1, 4] yang mewakili x1, y1, x2, y2 dari setiap kotak pembatas.[batch, num_boxes, num_classes] yang menunjukkan skor semua kelas untuk setiap kotak pembatas.Sampai sekarang, masih ada sepotong kecil pasca-pemrosesan termasuk NMS diperlukan. Kami berusaha meminimalkan waktu dan kompleksitas pasca-pemrosesan.
Script ini adalah untuk mengubah model darknet pretrained resmi menjadi onnx
Versi Pytorch direkomendasikan:
Instal onnxruntime
pip install onnxruntimeJalankan skrip Python untuk menghasilkan model ONNX dan jalankan demo
python demo_darknet2onnx.py < cfgFile > < namesFile > < weightFile > < imageFile > < batchSize > Anda dapat mengubah model Pytorch terlatih Anda menjadi ONNX menggunakan skrip ini
Versi Pytorch direkomendasikan:
Instal onnxruntime
pip install onnxruntimeJalankan skrip Python untuk menghasilkan model ONNX dan jalankan demo
python demo_pytorch2onnx.py < weight_file > < image_path > < batch_size > < n_classes > < IN_IMAGE_H > < IN_IMAGE_W >Misalnya:
python demo_pytorch2onnx.py yolov4.pth dog.jpg 8 80 416 416Jalankan perintah berikut untuk mengonversi model ONNX YOLOV4 menjadi mesin Tensorrt
trtexec --onnx= < onnx_file > --explicitBatch --saveEngine= < tensorRT_engine_file > --workspace= < size_in_megabytes > --fp16Jalankan perintah berikut untuk mengonversi model ONNX YOLOV4 menjadi mesin 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 > --fp16Misalnya:
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 >Demo ini di sini hanya berfungsi ketika BatchSize dinamis (1 harus berada dalam rentang dinamis) atau BatchSize = 1, tetapi Anda dapat memperbarui demo ini sedikit untuk ukuran batch dinamis atau statis lainnya.
Note1: input_h dan input_w harus setuju dengan ukuran input dalam file onnx asli.
Catatan2: Operasi NMS tambahan diperlukan untuk output TensorRT. Demo ini menggunakan kode NMS Python dari tool/utils.py .
Pertama: Konversi ke Onnx
TensorFlow> = 2.0
1: Terima kasih: github: https: //github.com/onnx/onnx-tensorflow
2: Jalankan git clone https://github.com/onnx/onnx-tensorflow.git && cd onnx-tensorflow run pip install -e.
Catatan: Kesalahan akan terjadi saat menggunakan "Pip Instal Onnx-TF", setidaknya untuk saya, disarankan untuk menggunakan instalasi kode sumber
cd DeepStream
make
Untuk batch tunggal,
trtexec --onnx=<onnx_file> --explicitBatch --saveEngine=<tensorRT_engine_file> --workspace=<size_in_megabytes> --fp16
Untuk multi-batch,
trtexec --onnx=<onnx_file> --explicitBatch --shapes=input:Xx3xHxW --optShapes=input:Xx3xHxW --maxShapes=input:Xx3xHxW --minShape=input:1x3xHxW --saveEngine=<tensorRT_engine_file> --fp16
Catatan: Maxshapes tidak bisa lebih besar dari bentuk model asli.
Referensi:
@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}
}
