onnxruntime: ONNX Runtime much slower than PyTorch (2-3x slower)

Describe the bug We built a UNet3D image segmentation model in PyTorch (based on this repo) and want to start distributing it. ONNX seemed like a good option as it allows us to compress our models and the dependencies needed to run them. As our models are large & slow, we need to run them on GPU

We were able to convert these models to ONNX, but noticed a significant slow-down of the inference (2-3x). The issue is that the timing is quite critical, and that our models are already relatively slow, so we can’t afford more slow-downs

I’m running my comparison tests following what was done in this issue

I could use your help to better understand where the issue is coming from and if it is resolvable at all. What tests, settings, etc. can I try to see where the issue might be ?

Urgency This is quite an urgent issue, we need to deliver our models to our clients in the coming month and will need to resolve to other solutions if we can’t fix ONNX soon

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 18.04
• ONNX Runtime installed from (source or binary): source
• ONNX Runtime version: 1.12
• Python version: 3.8.13
• Visual Studio version (if applicable):
• CUDA/cuDNN version: 11.1 / 8.0.5
• GPU model and memory: NVIDIA GeForce RTX 3060, 12GB

To Reproduce The code for the model will be quite hard to extract, so I’ll first try to describe the issue and what I’ve tested. I’m currently generating my model using:

with torch.no_grad():
        torch.onnx.export(
            torchModel,
            dummyInput,
            outPath,
            export_params=True,
            opset_version=14,
            do_constant_folding=True,
            input_names=["input"],
            output_names=["output"],
            dynamic_axes={
                "input": [0],
                "output": [0]
            },
            verbose=True,
        )

The model that we are using uses the following operations:

• 3D conv
• Group normalisation
• Max pooling
• Nearest neighbor interpolation & concatenation

When converting to ONNX, I could see some weird things in the graph (see the first screenshot):

1. For some operations the device is cpu instead of cuda:0 like all other operations; what does this mean? Will ONNX runtime run these operations on the CPU?. See below for the partial output of the conversion with verbose = True

  %154 : Half(*, 512, *, *, *, strides=[884736, 1728, 144, 12, 1], requires_grad=0, device=cuda:0) = onnx::Relu(%153) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:1297:0
  %155 : Long(3, strides=[1], device=cpu) = onnx::Constant[value= 0  8 -1 [ CPULongType{3} ]]()
  %156 : Half(0, 8, *, device=cpu) = onnx::Reshape[allowzero=0](%154, %155)
  %157 : Half(8, strides=[1], device=cpu) = onnx::Constant[value= 1  1  1  1  1  1  1  1 [ CPUHalfType{8} ]]()
  %158 : Half(8, strides=[1], device=cpu) = onnx::Constant[value= 0  0  0  0  0  0  0  0 [ CPUHalfType{8} ]]()
  %159 : Half(0, 8, *, device=cpu) = onnx::InstanceNormalization[epsilon=1.0000000000000001e-05](%156, %157, %158)
  %160 : Long(5, strides=[1], device=cpu) = onnx::Shape(%154)
  %161 : Half(*, *, *, *, *, device=cpu) = onnx::Reshape[allowzero=0](%159, %160)
  %164 : Half(*, *, *, *, *, device=cpu) = onnx::Mul(%161, %309)
  %167 : Half(*, *, *, *, *, strides=[884736, 1728, 144, 12, 1], requires_grad=0, device=cuda:0) = onnx::Add(%164, %310) # /usr/local/lib/python3.8/dist-packages/torch/nn/functional.py:2360:0
  %169 : Long(5, strides=[1], device=cpu) = onnx::Shape(%167)
  %170 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={0}]()
  %171 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={0}]()
  %172 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={2}]()
  %173 : Long(2, strides=[1], device=cpu) = onnx::Slice(%169, %171, %172, %170)
  %175 : Long(5, strides=[1], device=cpu) = onnx::Concat[axis=0](%173, %311)
  %176 : Tensor? = prim::Constant()
  %177 : Tensor? = prim::Constant()

1. The graph is rather “ugly” when compared to the one generated for ResNet (regarding all the Mul, Add, Reshape, etc. operations). Could this be the reason for the slowdown?

I saw that Group normalisation wasn’t directly supported by ONNX and thus thought that this might be the cause for the slow-down, I thus tried with an alternative model where I remove the group norm, which led to a nicer graph (see 2nd screenshot) and to less slow-down (from 3x slower to 2x slower). The slow-down is still significant though, and the Slice, Concat, etc. operations still say that they occur on the cpu; are these then the issue?

Overall it would be great to get some guidance on where the problem could be located: should we adapt our model architecture, the way of exporting to ONNX, etc. ? Is it even possible at all with a model like UNet3D ?

Thanks for the help !

Screenshots