Creating a TensorRT engine from an ONNX model file

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

I am attempting to use an onnx model to create a tensorrt engine using the python api. Currently, the onnx file is not properly parsed. I attempted using the last layer as output without success.

import os
import tensorrt as trt

TRT_LOGGER = trt.Logger()

def build_engine_onnx(model_file):
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
        builder.max_workspace_size = 1 << 20
        builder.max_batch_size = 1

        with open(model_file, 'rb') as model:
            parser.parse(model.read())

        network.mark_output(network.get_layer(network.num_layers - 1).get_output(0))

        return builder.build_cuda_engine(network)

engine = build_engine_onnx("frozen.onnx")

engine_file_path = "tensorrt.engine"
with open(engine_file_path, "wb") as f:
    f.write(engine.serialize())

After calling build_engine_onnx the engine only contains an input layer. Execution does not result in any errors. Is it possible to specify the output layer by name or display the network layer names with indices?

Any help is highly appreciated!

2

Hi,

Please refer to below sample example:

You can refer to “onnx_to_tensorrt.py” file to build a TensorRT engine from the ONNX file using python API.

Thanks