Iplugin tensorrt engine error for ds5.0

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.

1 Like