Conversion of FasterRCNN ONNX model to TensorRT fails

Description

I am trying to convert an ONNX FasterRCNN based object detection model to a TensorRT engine and running into an error with shapes in the ROI heads.

Here is an example script that will allow one to recreate the error:

import os import tensorrt as trt if __name__ == '__main__': x_shape = (2,2,300,400) DIR_NAME = os.path.dirname(__file__) TRT_LOGGER = trt.Logger(trt.Logger.VERBOSE) trt.init_libnvinfer_plugins(TRT_LOGGER, '') builder = trt.Builder(TRT_LOGGER) network = builder.create_network() parser = trt.OnnxParser(network, builder.logger) onnx_path = os.path.realpath('model_zoo_rcnn.onnx') with open(onnx_path, "rb") as f: if not parser.parse(f.read()): TRT_LOGGER.log(trt.Logger.ERROR, f"Failed to load ONNX file: {onnx_path}") for error in range(parser.num_errors): TRT_LOGGER.log( trt.Logger.ERROR, parser.get_error(error)) raise RuntimeError(f'Failed to load ONNX file: {onnx_path}') engine_path = os.path.realpath('model_zoo_rcnn.engine') engine_dir = os.path.dirname(engine_path) os.makedirs(engine_dir, exist_ok=True) inputs = [network.get_input(i) for i in range(network.num_inputs)] config = builder.create_builder_config() config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED engine_str = builder.build_serialized_network(network, config) with open(engine_path, "wb") as f: f.write(engine_str) 

And here is the model used for this example:

This occurs when using the default TensorRT plugin based ROIAlign node. It occurs even when an exported ONNX model includes a fixed batch size/input and output tensor shapes. These issues appear to occur with any RCNN based model, including ones from the ONNX Model Zoo.

Does the input ONNX file/graph need to be modified to use dynamic shapes in this node or set the node’s inputs/outputs to be fixed? If so, how can this be done? If this is not the issue, what is causing this error and how can it be fixed?

Environment

TensorRT Version: 10.5
GPU Type: Nvidia Jetson Orin
Nvidia Driver Version: 540.4.0
CUDA Version: 12.6
CUDNN Version:
Operating System + Version: Ubuntu 22.04.04
Python Version (if applicable): 3.10.15
Baremetal or Container (if container which image + tag): Baremetal

Relevant Files

Log from loading ONNX graph:
faster_rcnn_ex.log (1.5 MB)

1 Like

I see this previous post from May 2023. Is the ONNX to TRT converter fixed to support this file type?

I have been able to get this workaround to work. We can do this through our onnx_graphsurgeon tool by forcing the condition to be a constant. TRT is able to optimize out the conditional entirely in this case.Modification script:

import onnx_graphsurgeon as gs import numpy as np const_else = gs.Constant("const_else", np.array([0]).astype(np.bool_)) g = gs.import_onnx(onnx.load("fasterrcnn.onnx")) # Overwrite condition to a constant. for n in g.nodes: if n.op == "If": n.inputs = [const_else] onnx.save(gs.export_onnx(g), "fasterrcnn_constelse.onnx") 

While converting MaskRCNN ONNX to trt, I used the above workaround and and running into new error.

Here is an example code to recreate the error:

def simple_export(): """Simple export process""" try: # Load original model print("πŸ“¦ Loading Mask R-CNN model...") model = torchvision.models.detection.maskrcnn_resnet50_fpn( pretrained=False, num_classes=NUM_CLASSES ) model.load_state_dict(torch.load(MODEL_PTH_PATH, map_location="cuda")) model.to("cuda") model.eval() # # Extract just the backbone # print("πŸ”§ Extracting backbone...") # backbone_model = BackboneOnly(model) # backbone_model = backbone_model.cuda() # backbone_model.eval() # Test the model first print("πŸ§ͺ Testing model...") dummy_input = torch.randn(*INPUT_SHAPE).cuda() with torch.no_grad(): output = model(dummy_input) # print(f"βœ… Model output shape: {output.shape}") # Export to ONNX print("πŸ“€ Exporting backbone to ONNX...") torch.onnx.export( model, dummy_input, ONNX_EXPORT_PATH, export_params=True, opset_version=11, do_constant_folding=True, input_names=["input"], output_names=['boxes', 'labels', 'scores', 'masks'], verbose=False ) print("βœ… ONNX export successful") const_else = gs.Constant("const_else", np.array([0]).astype(np.bool_)) g = gs.import_onnx(onnx.load("maskrcnn_simple.onnx")) # Overwrite condition to a constant. for n in g.nodes: if n.op == "If": n.inputs = [const_else] onnx.save(gs.export_onnx(g), "maskrcnn_constelse.onnx") return True except Exception as e: print(f"❌ Export failed: {e}") return False def convert_to_tensorrt(): """Convert ONNX to TensorRT""" try: print("πŸ”§ Converting to TensorRT...") # Create builder logger = trt.Logger(trt.Logger.INFO) builder = trt.Builder(logger) network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser = trt.OnnxParser(network, logger) # Parse ONNX with open("maskrcnn_constelse.onnx", 'rb') as model_file: if not parser.parse(model_file.read()): for error in range(parser.num_errors): print(parser.get_error(error)) return False print("βœ… ONNX parsed successfully") # Build config config = builder.create_builder_config() if hasattr(config, 'set_memory_pool_limit'): config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30) # Set input shape profile = builder.create_optimization_profile() input_name = network.get_input(0).name profile.set_shape(input_name, INPUT_SHAPE, INPUT_SHAPE, INPUT_SHAPE) config.add_optimization_profile(profile) # Build engine print("βš™οΈ Building engine...") serialized_engine = builder.build_serialized_network(network, config) if serialized_engine: with open(TRT_ENGINE_PATH, 'wb') as f: f.write(serialized_engine) print(f"βœ… Engine saved to {TRT_ENGINE_PATH}") return True else: print("❌ Failed to build engine") return False except Exception as e: print(f"❌ TensorRT conversion failed: {e}") return False