pytorch YOLOv4
1.0.0
Yolov4의 최소 Pytorch 구현.
종이 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 데이터 세트의 경우 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 ...
...
...
기차
cfg.py에서 매개 변수를 설정할 수 있습니다.
python train.py -g [GPU_ID] -dir [Dataset direction] ...
Onnx 및 Tensorrt 모델은 Pytorch (Tianxiaomo) : pytorch-> onnx-> tensorrt에서 변환됩니다. 전환에 대한 자세한 내용은 다음 섹션을 참조하십시오.
| 모델 유형 | ap | 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 |
| 모델 유형 | ap | 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 과 같은 다른 입력 비율에 대해 입력 크기를 조정할 수 있습니다. 입력 크기가 커지면 더 작은 대상을 감지하는 데 도움이 될 수 있지만 속도가 느리고 GPU 메모리가 소진 될 수 있습니다.
height = 320 + 96 * n , n in { 0 , 1 , 2 , 3 , ...}
width = 320 + 96 * m , m in { 0 , 1 , 2 , 3 , ...}사전에 사기꾼 모델 및 DarkNet Weights를로드하여 추론을 수행합니다 (이미지 크기는 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] 입니다.[batch, num_boxes, num_classes] 인 경계 상자의 점수입니다.지금까지 여전히 NMS를 포함한 작은 사후 처리가 필요합니다. 우리는 후 처리의 시간과 복잡성을 최소화하려고 노력하고 있습니다.
이 스크립트는 공식 사전 상인 된 DarkNet 모델을 ONNX로 변환하는 것입니다.
Pytorch 버전 권장 :
onnxRuntime을 설치하십시오
pip install onnxruntimePython 스크립트를 실행하여 Onnx 모델을 생성하고 데모를 실행하십시오.
python demo_darknet2onnx.py < cfgFile > < namesFile > < weightFile > < imageFile > < batchSize > 이 스크립트를 사용하여 훈련 된 Pytorch 모델을 Onx로 변환 할 수 있습니다.
Pytorch 버전 권장 :
onnxRuntime을 설치하십시오
pip install onnxruntimePython 스크립트를 실행하여 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 >이 데모는 Batchsize가 동적 (1이 동적 범위 내에 있어야 함) 또는 Batchsize = 1 일 때만 작동하지만 다른 동적 또는 정적 배치 크기에 대해이 데모를 약간 업데이트 할 수 있습니다.
Note1 : Input_h 및 Input_w는 원래 Onnx 파일의 입력 크기에 동의해야합니다.
참고 2 : Tensorrt 출력에는 추가 NMS 작업이 필요합니다. 이 데모는 tool/utils.py 의 Python NMS 코드를 사용합니다.
첫째 : ONNX 로의 전환
텐서 플로우> = 2.0
1 : 감사합니다 : github : https : //github.com/onnx/onnx-tensorflow
2 : 실행 git 클론 https://github.com/onnx/onnx-tensorflow.git && cd onnx-tensorflow run pip install -e.
참고 : "PIP 설치 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
참고 : Max 셰이프는 모델 원래 모양보다 클 수 없습니다.
참조:
@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}
}
