ONNX support for OneFormer Semantic Segmentation Task

[SysDevHayes]

Feature request

I'm trying to export a pretrained OneFormer to ONNX. I know that optimum has not yet officially supported exporting OneFormer to ONNX. That was why I wrote my own export script. Here it is:

from transformers import OneFormerProcessor, OneFormerForUniversalSegmentation 
from PIL import Image
import requests
import torch

processor = OneFormerProcessor.from_pretrained("shi-labs/oneformer_coco_swin_large")
model = OneFormerForUniversalSegmentation.from_pretrained("shi-labs/oneformer_coco_swin_large")

inputs = processor(Image.open('./sample.jpg'), ["semantic"], return_tensors="pt")

torch.onnx.export(
    model,
    [inputs],
    "export.onnx",
    export_params=True,
    opset_version=16,
    do_constant_folding=True,
    input_names = ['pixel_values'],
    output_names = ['output'],
    dynamic_axes={'input' : {0 : 'batch_size'},
                  'output' : {0 : 'batch_size'}
                 }
)

This gave me the error:

File /media/research/venv_3.10_2/lib/python3.10/site-packages/torch/nn/modules/module.py:1508, in Module._slow_forward(self, *input, **kwargs)
   1506         recording_scopes = False
   1507 try:
-> 1508     result = self.forward(*input, **kwargs)
   1509 finally:
   1510     if recording_scopes:

TypeError: OneFormerForUniversalSegmentation.forward() missing 1 required positional argument: 'task_inputs'

Motivation

If OneFormer can be exported to ONNX, it can be used in TensorRT. Which would boost the speed greatly while, at the same time, reducing the amount of computing resources.

Your contribution

I came across this one. However, with that approach, the exported ONNX model would not work on GPU. Since, the implementation by SHILAB utilized MultiScaleDeformableAttention, which I believe would not work in torch.trace

I don't think Transformer is using MultiScaleDeformableAttention. Yet still, I'm quite new to this and I would really appreciate if someone can point me to the right direction.

Thank you so much!