ONNX to TensorRT with dynamic batch size in Python

AI & Data Science Deep Learning (Training & Inference) TensorRT
,
1

Description

Trying to convert the yolov3-tiny-416 model to TensorRT with a dynamic batch size, with code modified from tensorrt_demos/yolo at master · jkjung-avt/tensorrt_demos · GitHub

The resulting engine is always None. Code snippets below.

Environment

Using the docker container nvcr.io/nvidia/tensorrt:20.08-py3
TensorRT Version: 7.1.3.4
GPU Type: Titan X
Nvidia Driver Version: 450.51.06
CUDA Version: 11.0
CUDNN Version:
Operating System + Version: Ubuntu 18.04
Python Version (if applicable): 3.6.9
ONNX Version: 1.4.1
TensorFlow Version (if applicable):
PyTorch Version (if applicable):
Baremetal or Container (if container which image + tag): nvcr.io/nvidia/tensorrt:20.08-py3

Relevant Files

Modified the build_engine function in tensorrt_demos/onnx_to_tensorrt.py at master · jkjung-avt/tensorrt_demos · GitHub to

def build_engine(onnx_file_path, category_num=80, verbose=True):
"""Build a TensorRT engine from an ONNX file."""
TRT_LOGGER = trt.Logger(trt.Logger.VERBOSE) # if verbose else trt.Logger()
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(*EXPLICIT_BATCH) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
    builder.max_workspace_size = 1 << 30
    # builder.max_batch_size = 32
    builder.fp16_mode = True
    #builder.strict_type_constraints = True
    
    config = builder.create_builder_config()

    # Parse model file
    print('Loading ONNX file from path {}...'.format(onnx_file_path))
    with open(onnx_file_path, 'rb') as model:
        if not parser.parse(model.read()):
            print('ERROR: Failed to parse the ONNX file.')
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return None
    
	shape = list(network.get_input(0).shape)
    shape[0] = -1
    network.get_input(0).shape = shape
    print(network.get_input(0).shape)

    print('Adding yolo_layer plugins...')
    model_name = onnx_file_path[:-5]
    network = add_yolo_plugins(
        network, model_name, category_num, TRT_LOGGER)
    
    
    profile = builder.create_optimization_profile()
    profile.set_shape(network.get_input(0).name, (1, 3, 416, 416), (16, 3, 416, 416), (32, 3, 416, 416))
    config.add_optimization_profile(profile)

    print('Building an engine.  This would take a while...')
    print('(Use "--verbose" to enable verbose logging.)')
    engine = builder.build_engine(network, config)
    print('Completed creating engine.')
    return engine

ONNX model: https://drive.google.com/file/d/1-WJCijVL9_JdEVVLOanzGvl9RmIyJWF_/view?usp=sharing
Model IR version: 4
Opset Version: 9

Steps To Reproduce

  • Run the above build_engine function with the ONNX model in the link, the resulting engine is always None

I am not sure I am missing something simple, or if there is a compatibility issue here. I thought adding the optimization profile would do the trick.

Any help is much appreciated.

2

Hi @aravind.anantha,
I tried generating engine using trtexec command , and it worked fine for me.
Can you try the below command once.
trtexec --onnx=yolov3-tiny-416.onnx --verbose --explicitBatch --shapes=000_net:1x3x416x416

Thanks!

4

Hi @AakankshaS I saved the engine this way, and loaded it back with the Python API to check it.
engine.get_binding_shape(0) (-1, 1, 224, 224)

But, when I see engine.max_batch_size, it is 1

I’m not sure if I need to change anything else to make it work.

This is the command I used.

trtexec --onnx=yolov3-tiny-416.onnx --explicitBatch --optShapes=000_net:16x3x416x416 --maxShapes=000_net:32x3x416x416 --minShapes=000_net:1x3x416x416 --shapes=000_net:8x3x416x416 --saveEngine=yolov3-tiny-416.trt

5

I realized the difference between execute_async() and execute_async_v2(). The latter ignores the engine batch size and is used for dynamic batches.

I was able to run a Python script with the engine generated using trtexec command from the above comment with different batch sizes

1 Like