Home

Awesome

Pytorch-YOLOv4

A minimal PyTorch implementation of YOLOv4.

├── 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

image

0. Weights Download

0.1 darknet

0.2 pytorch

you can use darknet2pytorch to convert it yourself, or download my converted model.

1. Train

use yolov4 to train your own data

  1. Download weight

  2. Transform data

    For coco dataset,you can use 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 ...
    ...
    ...
    
  3. Train

    you can set parameters in cfg.py.

     python train.py -g [GPU_ID] -dir [Dataset direction] ...
    

2. Inference

2.1 Performance on MS COCO dataset (using pretrained DarknetWeights from https://github.com/AlexeyAB/darknet)

ONNX and TensorRT models are converted from Pytorch (TianXiaomo): Pytorch->ONNX->TensorRT. See following sections for more details of conversions.

Model typeAPAP50AP75APSAPMAPL
DarkNet (YOLOv4 paper)0.4710.7100.5100.2780.5250.636
Pytorch (TianXiaomo)0.4660.7040.5050.2670.5240.629
TensorRT FP32 + BatchedNMSPlugin0.4720.7080.5110.2730.5300.637
TensorRT FP16 + BatchedNMSPlugin0.4720.7080.5110.2730.5300.636
Model typeAPAP50AP75APSAPMAPL
DarkNet (YOLOv4 paper)0.4120.6280.4430.2040.4440.560
Pytorch (TianXiaomo)0.4040.6150.4360.1960.4380.552
TensorRT FP32 + BatchedNMSPlugin0.4120.6250.4450.2000.4460.564
TensorRT FP16 + BatchedNMSPlugin0.4120.6250.4450.2000.4460.563

2.2 Image input size for inference

Image input size is NOT restricted in 320 * 320, 416 * 416, 512 * 512 and 608 * 608. You can adjust your input sizes for a different input ratio, for example: 320 * 608. Larger input size could help detect smaller targets, but may be slower and GPU memory exhausting.

height = 320 + 96 * n, n in {0, 1, 2, 3, ...}
width  = 320 + 96 * m, m in {0, 1, 2, 3, ...}

2.3 Different inference options

2.4 Inference output

There are 2 inference outputs.

Until now, still a small piece of post-processing including NMS is required. We are trying to minimize time and complexity of post-processing.

3. Darknet2ONNX

3.1 Dynamic or static batch size

4. Pytorch2ONNX

4.1 Dynamic or static batch size

5. ONNX2TensorRT

5.1 Convert from ONNX of static Batch size

5.2 Convert from ONNX of dynamic Batch size

5.3 Run the demo

python demo_trt.py <tensorRT_engine_file> <input_image> <input_H> <input_W>

6. ONNX2Tensorflow

7. ONNX2TensorRT and DeepStream Inference

  1. Compile the DeepStream Nvinfer Plugin
    cd DeepStream
    make 
  1. Build a TRT Engine.

For single batch,

trtexec --onnx=<onnx_file> --explicitBatch --saveEngine=<tensorRT_engine_file> --workspace=<size_in_megabytes> --fp16

For multi-batch,

trtexec --onnx=<onnx_file> --explicitBatch --shapes=input:Xx3xHxW --optShapes=input:Xx3xHxW --maxShapes=input:Xx3xHxW --minShape=input:1x3xHxW --saveEngine=<tensorRT_engine_file> --fp16

Note :The maxShapes could not be larger than model original shape.

  1. Write the deepstream config file for the TRT Engine.

Reference:

@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}
}