Wrong Predictions PyTorch vs TensorRT

[dudeperf3ct]

Environment

TensorRT Version: 7.1.3.0
GPU Type:
Nvidia Driver Version:
CUDA Version: 10.2.89
CUDNN Version: 8.0
Operating System + Version: Jetson Xavier Ubuntu 18.04 + Jetpack Revision:4.4 BOARD: t168ref EABI: aarch64
Python Version (if applicable): 3.6
TensorFlow Version (if applicable):
PyTorch Version (if applicable): 1.7.0
Baremetal or Container (if container which image + tag):

Description

Output mismatch between PyTorch prediction and TensorRT predictions.

I have trained a semantic segmentation model consisting of mobilenetv2 encoder and aspp_residual decoder. The outputs from PyTorch don't match the ones we get from converting onnx model to tensorrt. PyTorch is producing correct predictions but TensorRT predictions are all same. I have done a comparision of pytorch and onnx outputs using onnx runtime for different tolerance values. They seem alright.

Relevant Files

Attaching predictions from PyTorch and TensorRT along with scripts, prediction images, and trained models to re-run the model. Also, attached the results of PyTorch and onnx prediction (for 3 tolerance values 1e-1, 1e-2, and 1e-3).

Drive Link: https://drive.google.com/drive/folders/1RG2kWEpROgFQI7YbQBig13Rozi41RHr0?usp=sharing

Steps To Reproduce

• PyTorch prediction: python infer_img -l logs_pytorch/ -p models -i carot_2074.jpg --backend pytorch
• TensorRT prediction: python infer_img -l logs_tensorrt/ -p models -i carot_2074.jpg --backend tensorrt

Comparison

PyTorch Prediction

Pytorch -- Onnx --- Ground Truth predictions (for 1e-1 tolerance)

TensorRT prediction

[dudeperf3ct]

The error occurs because of adding a argmax layer on top of onnx model and then attempting to convert it to tensorRT.

# architecture definition from onnx if no engine is there
# get weights?
if self.engine is None:
  try:
    # basic stuff for onnx parser
    self.model_path = path + "/model.onnx"
    self.network = self.builder.create_network(network_flags)
    self.onnxparser = trt.OnnxParser(self.network, TRT_LOGGER)
    self.model = open(self.model_path, 'rb')
    self.onnxparser.parse(self.model.read())
    print("Successfully ONNX weights from ", self.model_path)
    # now add the argmax op in gpu
    self.topk = self.network.add_topk(
        self.network.get_output(0), trt.TopKOperation.MAX, 1, 1)
    assert(self.topk != None)
    self.topk.get_output(1).name = 'out'
    self.network.unmark_output(self.network.get_output(0))
    self.network.mark_output(self.topk.get_output(1))
    print("Successfully added argmax layer")
  except Exception as e:
    print("Couldn't load ONNX network. Error: ", e)
    quit()

The other approach would be skipping this step of adding argmax layer and instead choose to do a argmax on tensorRT outputs. This works charmingly! But adding a argmax layer and marking it seems to be the source of the problem. Any ideas on what's causing this?

[thhung]

@dudeperf3ct You closed this issue. Did you solve this problem? Could you explain what happened? I am having the same issue. Many thanks.

[dudeperf3ct]

@thhung, I found the alternative approach worked that's why I closed this issue. But the issue of adding and marking a argmax layer on top of onnx still persists.

The other approach would be skipping this step of adding argmax layer and instead choose to do a argmax on tensorRT outputs.

[tejal567]

I am also having the same problem. I have an argmax layer at the end of my network. TensorRT output is blank in this case. A workaround that I have done for now is to manually writing a function for doing argmax and removing argmax layer from network.

[nvpohanh]

I don't know if this is the same issue, but I have seen similar issue caused by using INT64 torch.Tensor or nump.ndarray to get TRT's output indices, which are in INT32.

I am going to close this for now. Please feel free to reopen if the issue still exists. Thanks