Unable to inference with malformed ONNX graph after export

[addisonklinke]

The Problem

I successfully exported the PyTorch checkpoint for tf_efficientdet_lite0 to ONNX by following the official PyTorch tutorial. However, when I try to create an inference session using the newly created ONNX model, the graph around node Concat_42 appears to be malformed since the inputs do not have the same number of dimensions

ort_session = onnxruntime.InferenceSession(path)  <-- raised by this Python line
onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : 
FAIL : Load model from tf_efficientdet_lite0.onnx failed:Node (Concat_42) Op (Concat) 
[ShapeInferenceError] All inputs to Concat must have same rank

Has anyone else been able to successfully export and run inference with an ONNX model?

To Reproduce

Below is my complete code for initializing the model (which automatically fetches pretrained weights from torchhub) and exporting to ONNX. I pip installed onnx==1.8.0 and onnxruntime==1.6.0 in addition to the required effdet packages (i.e. torch==1.7.1 and timm==0.3.4)

import onnx
import onnxruntime
import torch
import torch.onnx
from effdet.config.model_config import get_efficientdet_config
from effdet.efficientdet import EfficientDet

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

print('Initializing model')
config = get_efficientdet_config('tf_efficientdet_lite0')
model = EfficientDet(config, pretrained_backbone=True)

print('Exporting ONNX')
model.eval()
dummy = torch.randn(1, 3, 512, 512)
torch_outs = model(dummy)
path = 'tf_efficientdet_lite0.onnx'
torch.onnx.export(
    model=model,
    args=dummy,
    f=path,
    opset_version=11)  # Default opset 9 does not support non-constant padding

print('Loading ONNX and testing inference')
onnx_model = onnx.load(path)
onnx.checker.check_model(onnx_model)
ort_session = onnxruntime.InferenceSession(path)
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(dummy)}
ort_outs = ort_session.run(None, ort_inputs)

Additional Context

Here is a Netron visualization of the ONNX graph with Concat_42 highlighted red. It is towards the very beginning, in what appears to be some mean subtraction and pixel normalization preprocessing steps. Note I enabled View > Show Attributes (ctrl + D) and View > Show Names (ctrl + U) before taking the screenshot

[addisonklinke]

After additional testing, this appears to be a compatibility issue between timm.models.layers.get_same_padding() and PyTorch's ONNX exporter. The rank mismatch for concat occurs in PyTorch 1.7.1 and 1.6.0, but not 1.5.0

Despite getting a valid ONNX graph in 1.5.0, there are still warnings regarding the trace and integer division. I think the trace warning is benign since the same padding is only calculated once to initialize the model and therefore should generalize to other inputs. However, the division warning seems like it could be responsible for the incompatibility with later versions of PyTorch. Any thoughts @rwightman?

Initializing model
Exporting ONNX
/opt/conda/conda-bld/pytorch_1587428270644/work/aten/src/ATen/native/BinaryOps.cpp:81: 
UserWarning: Integer division of tensors using div or / is deprecated, 
and in a future release div will perform true division as in Python 3. 
Use true_divide or floor_divide (// in Python) instead.

timm/models/layers/padding.py:19: TracerWarning: 
Converting a tensor to a Python float might cause the trace to be incorrect. 
We can't record the data flow of Python values, so this value will be treated as a constant in the future. 
This means that the trace might not generalize to other inputs!
  return max((math.ceil(x / s) - 1) * s + (k - 1) * d + 1 - x, 0)

[rwightman]

See #155 and linked issues, that padding issue is easy enough to work around, but it's the smallest one in path to ONNX export