Pytorch to ONNX to TensorRT, segmentation channels output map incorrect when more than 1

I am trying to convert a segmentation model that is much like unet from pytorch to tensorrt. When I have 1 input channel and one 1 output channel it works correctly.

When I make the output channels of the Pytorch model 3 and do the conversion I get a 3x3 copy of the output map that is 3 channels deep in tensorrt. The top row of the 3x3 is the “red” channel, the middle row is the “green” and the bottom row is the “blue”. Eeach column and channel are repeats of the data.

When I make the output 6 channels I get a grid that is 6x6 with same results as above. 6 rows of data in the image size, 6 repeat columns of that data and 6 channels that are all copies of this data.

In every case I am using gray scale images in so the input channel is 1

I fixed the problem here. The workflow I was using did not the channels in the correct location. After transposing to NCHW it was good to go. Plus some other minor details but the issue was all in my software, thanks for looking.

I had a similar problem trying to decode the output of my PyTorch/ONNX/TRT model. Do you happen to have your code publicly available so that I could take a look at how you handled it? Thanks.

Unfortunately the code belongs to the corporation I work for so I can’t share it.

There are a couple of things that helped me fixed my output that were oversights. First was the Pytorch network was trained on tensor values that were between 0-1 so I had to divide the input to the Tensorrt model by 255.

Second most importantly I didn’t have the input format in the NCHW format. Once I fixed these two things it was working much better.

One final detail is that my pytorch network output only a mask that binary 0 or 1. Now the tensorrt model somehow lost the threshold and I get the ‘analog’ values out so I am having to post process them with a threshold.

If you some specific error or description of your issue I am glad to try to advise. Is yours a segmentation model as well?

I would be more than happy to share my source code with you if you had the time to give it a glance?