Having trouble exporting model to ONNX

[JulianAndhika]

Hey, thanks for the great model, but i seem to be having some issues with exporting to ONNX.

I wanted to create a model that could perform faster and (possibly) lighter to run on a server. So far i've been using a docTR model that i am fine tuning for text recognition using my own dataset. I've been using pytorch, since i've been having some issues with tensorflow. I am now trying to export the model to ONNX, but i've been running to an issue following the documentation

Here's the code

import os
import torch
import requests
from flask import Flask, request, jsonify
from doctr.io import DocumentFile
from doctr.models import ocr_predictor
from doctr.models.recognition import crnn_mobilenet_v3_large
from doctr.datasets import VOCABS
from doctr.models.utils import export_model_to_onnx

os.environ["USE_TORCH"] = "1"
batch_size = 16
input_shape = (3,32,128)
reco_model = crnn_mobilenet_v3_large(pretrained=False, pretrained_backbone=False, vocab=VOCABS["legacy_french"])
reco_params = torch.load('weights/crnn_mobilenet_v3_large_20241206-105210.pt', map_location='cpu')
reco_model.load_state_dict(reco_params)
dummy_input = torch.rand((batch_size,*input_shape), dtype=torch.float32)
model_path = export_model_to_onnx(reco_model, model_name="crnn.onnx",dummy_input=dummy_input)

and here's the error

TracerWarning: Converting a tensor to a Python list 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 list(zip(words, probs.tolist()))
Traceback (most recent call last):
  File "c:\Users\USER\Documents\Kerja\MT\document_ocr\test.py", line 21, in <module>
    model_path = export_model_to_onnx(reco_model, model_name="crnn.onnx",dummy_input=dummy_input)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\USER\Documents\Kerja\MT\document_ocr\.venv\Lib\site-packages\doctr\models\utils\pytorch.py", line 158, in export_model_to_onnx 
    torch.onnx.export(
  File "C:\Users\USER\Documents\Kerja\MT\document_ocr\.venv\Lib\site-packages\torch\onnx\__init__.py", line 375, in export
    export(
  File "C:\Users\USER\Documents\Kerja\MT\document_ocr\.venv\Lib\site-packages\torch\onnx\utils.py", line 502, in export
    _export(
  File "C:\Users\USER\Documents\Kerja\MT\document_ocr\.venv\Lib\site-packages\torch\onnx\utils.py", line 1564, in _export
    graph, params_dict, torch_out = _model_to_graph(
                                    ^^^^^^^^^^^^^^^^
  File "C:\Users\USER\Documents\Kerja\MT\document_ocr\.venv\Lib\site-packages\torch\onnx\utils.py", line 1113, in _model_to_graph
    graph, params, torch_out, module = _create_jit_graph(model, args)
                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\USER\Documents\Kerja\MT\document_ocr\.venv\Lib\site-packages\torch\onnx\utils.py", line 997, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\USER\Documents\Kerja\MT\document_ocr\.venv\Lib\site-packages\torch\onnx\utils.py", line 904, in _trace_and_get_graph_from_model
    trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(
                                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\USER\Documents\Kerja\MT\document_ocr\.venv\Lib\site-packages\torch\jit\_trace.py", line 1500, in _get_trace_graph
    outs = ONNXTracedModule(
           ^^^^^^^^^^^^^^^^^
  File "C:\Users\USER\Documents\Kerja\MT\document_ocr\.venv\Lib\site-packages\torch\nn\modules\module.py", line 1736, in _wrapped_call_impl     
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\USER\Documents\Kerja\MT\document_ocr\.venv\Lib\site-packages\torch\nn\modules\module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\USER\Documents\Kerja\MT\document_ocr\.venv\Lib\site-packages\torch\jit\_trace.py", line 139, in forward
    graph, out = torch._C._create_graph_by_tracing(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: output 1 (0.807211
[ CPUDoubleType{} ]) of traced region did not have observable data dependence with trace inputs; this probably indicates your program cannot be understood by the tracer.

Any ideas on what the issue is?

[felixdittrich92]

Hi @JulianAndhika 👋,

You missed to set the model to exportable :)

reco_model = crnn_mobilenet_v3_large(pretrained=False, pretrained_backbone=False, vocab=VOCABS["legacy_french"], exportable=True)

Best,
Felix

[felixdittrich92]

And you should set the batch size to None that it can handle any batch size (dynamic axis) :)

[JulianAndhika]

I forgot to use exportable lol. Thanks. Also,

And you should set the batch size to None that it can handle any batch size (dynamic axis) :)

Do i just set the batch size = None?

Running it with a batch size = 1 is successful, but running it with a batch size = None gives me an error of

Traceback (most recent call last):
  File "c:\Users\USER\Documents\Kerja\MT\document_ocr\test.py", line 20, in <module>
    dummy_input = torch.rand((batch_size,*input_shape), dtype=torch.float32)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: rand() received an invalid combination of arguments - got (tuple, dtype=torch.dtype), but expected one of:
 * (tuple of ints size, *, torch.Generator generator, tuple of names names, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
 * (tuple of ints size, *, torch.Generator generator, Tensor out = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
 * (tuple of ints size, *, Tensor out = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
 * (tuple of ints size, *, tuple of names names, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)

Also different than the documentations im using *input shape on dummy_input = torch.rand((batch_size,*input_shape), dtype=torch.float32), since not using * gives me an error. Is it correct to do it like that?

[felixdittrich92]

@JulianAndhika Oh i see the missing unpacking * is not correct in the documentation 👍 So

dummy_input = torch.rand((batch_size, *input_shape), dtype=torch.float32)

is correct :)

Some models support None as batch size input but yeah you are right 1 is fine for all crnn_ models because the Decoder LSTM doesn't allows a dynamic bs while exporting.

By using then https://github.com/felixdittrich92/OnnxTR with your exported model both cases are handled under the hood

[JulianAndhika]

@JulianAndhika Oh i see the missing unpacking * is not correct in the documentation 👍 So

dummy_input = torch.rand((batch_size, *input_shape), dtype=torch.float32)

is correct :)

Some models support None as batch size input but yeah you are right 1 is fine for all crnn_ models because the Decoder LSTM doesn't allows a dynamic bs while exporting.

By using then https://github.com/felixdittrich92/OnnxTR with your exported model both cases are handled under the hood

Great, the program works well. Thanks for the fast and helpful reply :)