[models] Ensure all PyTorch models are ONNX exportable

[fg-mindee]

Most users of the library are more interested in existing pretrained models to use for inference rather than training. For this reason, it's important to ensure we can easily export those trained models.

• General

  • Add a unittest/CI job to check that all models can be exported with ONNX ([ONNX] Part1: update current onnx export  #976, [ONNX/Tests] Part2 add ONNX tests for TF and PT detection models #979, [ONNX/Tests] Part2 add ONNX tests for TF and PT recognition models  #980 )

• Classification

  • VGG ([onnx] classification models export  #830)
  • ResNet ([onnx] classification models export  #830)
  • MobileNet V3 ([onnx] classification models export  #830)
  • MAGC ResNet ([onnx] classification models export  #830)

• Text Detection

  • LinkNet [ONNX/Tests] Part2 add ONNX tests for TF and PT detection models #979
  • DBNet [ONNX/Tests] Part2 add ONNX tests for TF and PT detection models #979

• Text Recognition

  • CRNN [ONNX/Tests] Part2 add ONNX tests for TF and PT recognition models  #980
  • SAR [ONNX/Tests] Part2 add ONNX tests for TF and PT recognition models  #980
  • MASTER [ONNX/Tests] Part2 add ONNX tests for TF and PT recognition models  #980 ([FIX] MASTER AMP and onnxruntime issue with master PT #986)

  NOTE: master is exportable but in fact of an onnxruntime internal issue (Could not find an implementation for Trilu node - ORT Model Optimizer microsoft/onnxruntime#10994) the loading fails

[felixdittrich92]

@charlesmindee classification complete 👍

[frytoli]

I'm relatively green when it comes to ML engineering, but I've been attempting to look at exporting the DBNet models to ONNX. I think we're relatively close already and just need to convert inputs/outputs to a data structure compatible with ONNX (specifically not numpy.ndarray). Here's some sample code and the resulting error. I'm slowly working on solving this on my end, but I wanted to post it here in case this is a really easy problem to solve for someone else more familiar with the project.

import torch
from doctr.models.detection import db_resnet50_rotation

model = db_resnet50_rotation(pretrained=True)
torch.onnx.export(model, torch.randn(1, 3, 1024, 1024), "model.onnx", input_names=['input'], output_names=['output'], opset_version=13)

RuntimeError: Only tuples, lists and Variables are supported as JIT inputs/outputs. Dictionaries and strings are also accepted, but their usage is not recommended. Here, received an input of unsupported type: numpy.ndarray

[felixdittrich92]

Hi @frytoli 👋,

thanks for working on this 🤗

To clearify where the problem is:
If you run the following with verbose=True you get a more detailed output:

import torch
from doctr.models.detection import db_resnet50_rotation

model = db_resnet50_rotation(pretrained=True)
dummy_input = torch.rand((1, 3, 1024, 1024), dtype=torch.float32)
torch.onnx.export(
        model,
        dummy_input,
        "test.onnx",
        input_names=['input'],
        output_names=['logits'],
        dynamic_axes={'input': {0: 'batch_size'}, 'logits': {0: 'batch_size'}},
        export_params=True, opset_version=13, verbose=True
    )

problem: preds[0] for preds in self.postprocessor(prob_map.detach().cpu().permute((0, 2, 3, 1)).numpy())

Currently the models forward pass includes the postprocessing call.
What in my opinion is not optimal to avoid something like this, when modeling, it should be noted that the model only returns specific data, i.e. logits vector and, if ground truth is present, loss.
So what should be done in the first place is to make these changes for all detection and recognition models.
However, doing so would break the pure pipeline, so we would have to see how it affects runtime/latency.
Furthermore, we would have to extend these changes to the Tensorflow part as well, in order to maintain equality on the one hand and on the other hand I think that we can then export the models more easily in Tensorflow's SavedModel format.
If you are interested in it, please open a pull request and we can work on it together if you like.
@charlesmindee @frgfm what is your opinion about this ?

[charlesmindee]

Hi @felixdittrich92 and @frytoli,
Thanks for bringing this on the table, to avoid breaking the whole pipeline as Felix mentioned it, we may want to create a flag ("exportable" or "raw" or something like that) that could be passed in the model initialization: db_resnet50(pretrained=True, exportable=True), and this would return in the call of the model the raw output without the postprocessing, or with an adapted post processing (ONNX compatible), what do you think ?
@MaximeChurin

[felixdittrich92]

@charlesmindee @frytoli Sounds good to me, but we have to take care of increasing complexity and the differences between TF and PT implementations. So i would suggest maybe to try it with on model but on both implementations so that we can provide an equivalent implementation which solves the ONNX export problem on PT and the SavedModel export on TF side. Wdyt ?

[charlesmindee]

Yes it would be nice to do a POC on 1 model first, but we may want to make it exportable in ONNX for both pytorch and tensorflow

[felixdittrich92]

So maybe if we can handle both with Onnx do we really need TF's SavedModel also ? Otherwise i would say let us modify this issue and close the counterpart issue for TF SavedModel 👍 (I will test tf2onnx this week with the classification models)