Plugin fails when trying YOLOv4-tiny model

Hey Alexis. Thank you for the great work.

I tried to run YOLOv4-tiny-416 model using an ONNX file converted using https://github.com/jkjung-avt/tensorrt_demos/blob/master/yolo/yolo_to_onnx.py.

However, the engine is not built, since an assertion in the YOLO layer plugin fails. python: yolo_layer.cu:115: virtual nvinfer1::Dims nvinfer1::YoloLayerPlugin::getOutputDimensions(int, const nvinfer1::Dims*, int): Assertion `inputs[0].d[0] == (mNumClasses + 5) * mNumAnchors' failed.

This is the class that I have added to yolo.py

class YOLOv4_tiny(YOLO):
    ENGINE_PATH = Path(__file__).parent / 'yolov4-tiny-416.trt'
    MODEL_PATH = Path(__file__).parent /  'yolov4-tiny-416.onnx'
    NUM_CLASSES = 90
    INPUT_SHAPE = (3, 416, 416)
    LAYER_FACTORS = [32, 16]
    SCALES = [1.05, 1.05]
    ANCHORS = [[81,82,  135,169,  344,319],
              [10,14,  23,27,  37,58]]

Could you please help me debug this?

[PiyalGeorge]

@GeekAlexis Thanks a lot. Will try that.

I tried your repo in Xavier-NX. It works like charm with Good FPS and accuracy! Good job

[GeekAlexis]

@PiyalGeorge Yes, you need to train a ReID model on the same class, assuming the class is not person. And there is nothing specific for tiny because ReID is independent of detector. FYI, you shouldn’t convert YOLO to TensorRT, converting to ONNX is enough.