pytorch YOLOv4
1.0.0
Минимальная реализация питтера Yolov4.
Paper Yolo V4: https://arxiv.org/abs/2004.10934
Исходный код: https: //github.com/alexeyab/darknet
Более подробная информация: http://pjreddie.com/darknet/yolo/
Вывод
Тренироваться
├── 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
Вы можете использовать Darknet2pytorch для преобразования его самостоятельно или загрузить мою конвертированную модель.
Используйте Yolov4 для обучения собственных данных
Скачать вес
Преобразовать данные
Для набора данных Coco вы можете использовать инструмент/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 ...
...
...
Тренироваться
Вы можете установить параметры в cfg.py.
python train.py -g [GPU_ID] -dir [Dataset direction] ...
Модели ONNX и Tensorrt преобразуются из Pytorch (Tianxiaomo): pytorch-> onnx-> tensorrt. См. Следующие разделы для более подробной информации о преобразовании.
| Тип модели | Доступа | AP50 | AP75 | APS | APM | Apl |
|---|---|---|---|---|---|---|
| Darknet (Yolov4 Paper) | 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 |
| Тип модели | Доступа | AP50 | AP75 | APS | APM | Apl |
|---|---|---|---|---|---|---|
| Darknet (Yolov4 Paper) | 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 |
Размер входа изображения не ограничен в 320 * 320 , 416 * 416 , 512 * 512 и 608 * 608 . Вы можете настроить размеры входа для другого входного отношения, например: 320 * 608 . Больший размер входа может помочь обнаружить меньшие цели, но может быть медленнее, а память графического процессора исчерпана.
height = 320 + 96 * n , n in { 0 , 1 , 2 , 3 , ...}
width = 320 + 96 * m , m in { 0 , 1 , 2 , 3 , ...}Загрузите предварительную модель Darknet и веса Darknet, чтобы сделать вывод (размер изображения уже настроен в файле CFG)
python demo.py -cfgfile < cfgFile > -weightfile < weightFile > -imgfile < imgFile >Загрузить веса Pytorch (PTH -файл), чтобы сделать вывод
python models.py < num_classes > < weightfile > < imgfile > < IN_IMAGE_H > < IN_IMAGE_W > < namefile(optional) >Нагрузка преобразована ONNX -файл для вывода (см. Раздел 3 и 4)
Нагрузка конвертированного файла двигателя Tensorrt для вывода (см. Раздел 5)
Есть 2 вывода вывода.
[batch, num_boxes, 1, 4] , которая представляет x1, y1, x2, y2 каждой ограничивающей коробки.[batch, num_boxes, num_classes] указывающие оценки всех классов для каждой ограничивающей коробки.До сих пор требуется небольшая часть пост-обработки, включая NMS. Мы пытаемся минимизировать время и сложность пост-обработки.
Этот скрипт должен преобразовать официальную предварительную модель Darknet в ONNX
Рекомендуется версия Pytorch:
Установите onnxruntime
pip install onnxruntimeЗапустите скрипт Python для создания модели ONNX и запустите демо.
python demo_darknet2onnx.py < cfgFile > < namesFile > < weightFile > < imageFile > < batchSize > Вы можете преобразовать обученную модель Pytorch в ONNX, используя этот скрипт
Рекомендуется версия Pytorch:
Установите onnxruntime
pip install onnxruntimeЗапустите скрипт Python для создания модели ONNX и запустите демо.
python demo_pytorch2onnx.py < weight_file > < image_path > < batch_size > < n_classes > < IN_IMAGE_H > < IN_IMAGE_W >Например:
python demo_pytorch2onnx.py yolov4.pth dog.jpg 8 80 416 416Запустите следующую команду, чтобы преобразовать модель yolov4 onnx в двигатель Tensorrt
trtexec --onnx= < onnx_file > --explicitBatch --saveEngine= < tensorRT_engine_file > --workspace= < size_in_megabytes > --fp16Запустите следующую команду, чтобы преобразовать модель yolov4 onnx в двигатель 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 > --fp16Например:
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 >Эта демонстрация здесь работает только тогда, когда пакетный размер является динамическим (1 должен быть в динамическом диапазоне) или пакетный размер = 1, но вы можете немного обновить эту демонстрацию для других динамических или статических партий.
Примечание1: input_h и input_w должны согласоваться с размером ввода в исходном файле ONNX.
Примечание2: Дополнительные операции NMS необходимы для вывода Tensorrt. Эта демонстрация использует код Python NMS из tool/utils.py .
Сначала: преобразование в ONNX
Tensorflow> = 2,0
1: Спасибо: GitHub: https: //github.com/onnx/onnx-tensorflow
2: Запустите Git Clone https://github.com/onnx/onnx-tensorflow.git && cd onnx-tensorflow run pip install -e.
Примечание: ошибки будут возникнуть при использовании «PIP Install ONNX-TF», по крайней мере для меня, рекомендуется использовать установку исходного кода
cd DeepStream
make
Для одиночной партии,
trtexec --onnx=<onnx_file> --explicitBatch --saveEngine=<tensorRT_engine_file> --workspace=<size_in_megabytes> --fp16
Для многогранны,
trtexec --onnx=<onnx_file> --explicitBatch --shapes=input:Xx3xHxW --optShapes=input:Xx3xHxW --maxShapes=input:Xx3xHxW --minShape=input:1x3xHxW --saveEngine=<tensorRT_engine_file> --fp16
Примечание. Maxshapes не может быть больше, чем оригинальная модель.
Ссылка:
@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}
}
