Trouble with TensorRT inference pipeline for Segmentation on Jetson Orin Nano

Robotics & Edge Computing Jetson Systems Jetson Orin Nano
1

Hello,
I am using a Jetson Orin Nano Dev Kit and TensorRT to perform embedded deep learning segmentation in a quantized and offline scenario. My issue is that while the tensorRT engine builds and the inference script I wrote uses the engine without any errors, my inference is very wrong.

I am expecting output arrays of (1,1024) for input images of size (1,1024,1024) size. When I create a figure that stacks all output prediction arrays on top of each other, I see that practically every prediction is the same with almost imperceptible change, even though the images are very different. When I run regular inference using pytorch, this stacked array image looks heterogeneous, which is more believable.

Can you help me figure out why the predictions look nearly identical from image to image? I am not sure where the error in prediction is happening and I have tried things like checking precision levels for input data and trt engine expectations, validating ONNX file functionality, and utilizing example scripts on the TensorRT github repository. I am using a virtual environment and including TensorRT version 10.3 on a Jetpack 6.2 SDK.

Any help would be appreciated. I am happy to provide code if requested as well.

Thanks

3

Hi,

Do you mean the output dimension is not expected?
If so, could you run the engine with trtexec?

$ /usr/src/tensorrt/bin/trtexec --onnx=[model]

Moreover, please verify the ONNX model with ONNXRuntime first.
In general, we expect the same output between ONNXRuntime and TensorRT.

Thanks.

5

Hi,

You can give it a quick try to verify the engine.

For example, with mnist.onnx.
We first convert the model into a TensorRT engine with below command:

$ /usr/src/tensorrt/bin/trtexec --onnx=/usr/src/tensorrt/data/mnist/mnist.onnx --saveEngine=mnist.engine

Then run it with random input and print out the output tensor:

$ /usr/src/tensorrt/bin/trtexec --loadEngine=mnist.engine --dumpOutput
...
[06/05/2025-07:14:26] [I] Output Tensors:
[06/05/2025-07:14:26] [I] Plus214_Output_0: (1x10)
[06/05/2025-07:14:26] [I] -1.60912 -0.901602 1.55434 1.57656 -0.0528268 -0.897766 0.831164 0.671341 -0.281258 -0.289107
&&&& PASSED TensorRT.trtexec [TensorRT v100300] # /usr/src/tensorrt/bin/trtexec --loadEngine=mnist.engine --dumpOutput

Please compare the result to the TensorRT and ONNXRuntime to see if the value is reasonable. (in the expected range)

If so, please check the below sample for the binding buffer.

If not, please share your ONNX model with us so we can check it further.

Thanks.

7

Hi,
My apologies, I thought I had published my edit of the previous message.

I thought the ONNXRuntime was outputting proper masks because it was completing without any failures or warnings, but when I plotted the output mask arrays in a similar format to how I was plotting trt engine-based inference masks, I encountered the same issue. So i think the issue is with the ONNX conversion from pth weights and model. I am just following the instructions provided by the TensorRT GitHub page for converting my torch model to onnx format, but I am unsure why it is producing such uniform and incorrect results.

My onnx_exporter code is below:

import torch
import onnx
from custom_model import my_model
import sys

device = torch.device('cuda') # if torch.cuda.is_available() else 'cpu')
model = my_model(spatial_dims=2, in_channels=1, out_channels=3).to(device)

model.load_state_dict(torch.load(r"myweights.pth", map_location=device))

model.eval()

dummy_input = torch.rand(1, 1, 1024, 1024).to(device)

torch.onnx.export(model, dummy_input, 'new_model.onnx', opset_version=12, do_constant_folding=True, input_names=['input'], output_names=['output'], dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}})

print('ONNX model created')

onnx_model = onnx.load('new_model.onnx')
onnx.checker.check_model(onnx_model)

And my ONNXRuntime code takes the ONNX inference and convert my output (1,3,1024) array to an image which is a tiled version of the 1x1024 argmax output (to create a 1024x1024 image)

8

There is no update from you for a period, assuming this is not an issue anymore.
Hence, we are closing this topic. If need further support, please open a new one.
Thanks ~0618

Hi,

When you export the model to ONNX, do you see any error or warning shown on the console?

Thanks.