There is a good way to generate tensorRT YOLOv4 engine without needs of TRT plugins
Step1: Download pretrained YOLOv4 model
Model definition can be downloaded from here
https://raw.githubusercontent.com/AlexeyAB/darknet/master/cfg/yolov4.cfg
Pretrained weights can be downloaded from here
https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.weights
- Specify input image size
Open file yolov4.cfg
and set values of hight and width at header part of the cfg file
- Input size options of YoloV4 for inference
Input size | Output 1 | Output 2 | Output 3 | |
---|---|---|---|---|
Size option 1 | 3x608x608 | 255x76x76 | 255x38x38 | 255x19x19 |
Size option 2 | 3x512x512 | 255x64x64 | 255x32x32 | 255x16x16 |
Size option 3 | 3x416x416 | 255x52x52 | 255x26x26 | 255x13x13 |
Size option 4 | 3x320x320 | 255x40x40 | 255x20x20 | 255x10x10 |
Hint: hight and width of input could be different, for example, 608x416, 320x608
Step2: Download a GitHub repository that can help you convert YOLOv4 from darknet to pytorch
git clone https://github.com/Tianxiaomo/pytorch-YOLOv4.git
Step3: Generate onnx model
Create a python file darknet2onnx.py
with code as follows, and copy this python script into pytorch-YOLOv4
, and then execute it this way:
python darknet2onnx.py <cfgFile> <weightFile> <batchSize>
import sys
import torch
from tool.darknet2pytorch import Darknet
def fransform_to_onnx(cfgfile, weightfile, batch_size):
model = Darknet(cfgfile)
model.print_network()
model.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
# model.cuda()
x = torch.randn((batch_size, 3, model.height, model.width), requires_grad=True) #.cuda()
# Export the model
print('Export the onnx model ...')
torch.onnx.export(model,
x,
"yolov4_{}_3_{}_{}.onnx".format(batch_size, model.height, model.width),
export_params=True,
opset_version=11,
do_constant_folding=True,
input_names=['input'], output_names=['output_1', 'output_2', 'output_3'],
dynamic_axes=None)
print('Onnx model exporting done')
if __name__ == '__main__':
if len(sys.argv) == 4:
cfgfile = sys.argv[1]
weightfile = sys.argv[2]
batch_size = int(sys.argv[3])
fransform_to_onnx(cfgfile, weightfile, batch_size)
else:
print('Please execute this script this way:\n')
print('python darknet2onnx.py <cfgFile> <weightFile> <batchSize>')
Step4: Transform onnx model into TensorRT model
- Generate TensorRT engine in fp16 mode:
./trtexec --onnx=<onnx_file> --workspace=4096 --saveEngine=<engine_file> --fp16 --explicitBatch
- Generate TensorRT engine in int8 mode:
./trtexec --onnx=<onnx_file> --workspace=4096 --saveEngine=<engine_file> --int8 --explicitBatch
Step5: How to train YOLOv4 by your own dataset
The public pretrained YOLOv4 model may not be usefull in specified industrial scenarios.
If you want to re-train this model via Pytorch with your own dataset, just comment or remove the
line model.load_weights(...)
from the source, and write your own training code.