Onnx export wrong shape

[ldt-becker]

First of all, i am amazed by the output quality of doctr out of the box. I still had to fine tune it to support some quircks of our expected data.

When I load the exported fine tuned modell in doctr i get the following error:

  File "/home/ldt/ocr/train/test/bench_onnx.py", line 14, in <module>
    res = predictor(single_img_doc)
          ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/train-doctr-new/lib/python3.11/site-packages/onnxtr/models/predictor/predictor.py", line 132, in __call__
    word_preds = self.reco_predictor([crop for page_crops in crops for crop in page_crops], **kwargs)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/train-doctr-new/lib/python3.11/site-packages/onnxtr/models/recognition/predictor/base.py", line 70, in __call__
    raw = [self.model(batch, **kwargs)["preds"] for batch in processed_batches]
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/train-doctr-new/lib/python3.11/site-packages/onnxtr/models/recognition/predictor/base.py", line 70, in <listcomp>
    raw = [self.model(batch, **kwargs)["preds"] for batch in processed_batches]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/train-doctr-new/lib/python3.11/site-packages/onnxtr/models/recognition/models/crnn.py", line 137, in __call__
    logits = self.run(x)
             ^^^^^^^^^^^
  File "/home/user/miniconda3/envs/train-doctr-new/lib/python3.11/site-packages/onnxtr/models/engine.py", line 139, in run
    logits = self.runtime.run(self.output_name, {self.runtime_inputs.name: inputs}, run_options=run_options)[0]
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/train-doctr-new/lib/python3.11/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 321, in run
    return self._sess.run(output_names, input_feed, run_options)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxruntime.capi.onnxruntime_pybind11_state.RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running Reshape node. Name:'node_Reshape_213' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/tensor/reshape_helper.h:45 onnxruntime::ReshapeHelper::ReshapeHelper(const onnxruntime::TensorShape&, onnxruntime::TensorShapeVector&, bool) input_shape_size == size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{32,352,2,128}, requested shape:{32,1,256}

The export code looks like this:

batch_size = 1
input_shape = (3, 32, 128)
reco_model = crnn_vgg16_bn(pretrained=False, pretrained_backbone=False, exportable=True, vocab=VOCABS["german"])
reco_model.from_pretrained(f'/home/user/ocr/train/doctr/references/recognition/exp7.pt', map_location="cpu")
reco_model.eval()
dummy_input = torch.rand((batch_size, *input_shape), dtype=torch.float32)
model_path = export_model_to_onnx(
    reco_model,
    model_name="model",
    dummy_input=dummy_input
)

And the doctr code looks like this:

from doctr.models import ocr_predictor
from doctr.io import DocumentFile, Document
from doctr.models import ocr_predictor, from_hub, crnn_vgg16_bn
from doctr.models.recognition.vitstr.pytorch import vitstr_small
from doctr.models.utils import export_model_to_onnx
import torch
import sys
from doctr.datasets.vocabs import VOCABS

# Load your docTR model
batch_size = 1
input_shape = (3, 32, 128)
reco_model = crnn_vgg16_bn(pretrained=False, pretrained_backbone=False, exportable=True, vocab=VOCABS["german"])
reco_model.from_pretrained(f'/home/ldt/ocr/train/doctr/references/recognition/exp{sys.argv[1]}.pt', map_location="cpu")
reco_model.eval()
dummy_input = torch.rand((batch_size, *input_shape), dtype=torch.float32)
model_path = export_model_to_onnx(
    reco_model,
    model_name="model",
    dummy_input=dummy_input
)

I'd be very thankful for any help, thanks!