Add support for ORTModelForObjectDetection

[shivalikasingh95]

Feature request

Hi, I went through optimum's code base and could not find support for object detection models. Is there plan to add ORTModelForObjectDetection just like ORTModelForImageClassification exists? Would be great to have this feature.

Object detection task is also supported as part of transformers pipeline feature so I guess it should be possible to support this as part of optimum?

Motivation

I want to leverage onnx support for YOLOS model

Your contribution

I would be happy to help in adding support for this feature if someone can guide me.

[mht-sharma]

We are currently working on adding ONNX export features in optimum, and could add support to this architecture as well.

[shivalikasingh95]

Sounds great, thanks @mht-sharma! Looking forward to have this available!

[fxmarty]

Hello @shivalikasingh95 , great suggestion! If you want to contribute to adding ORTModelForObjectDetection, I suggest:

• Reading https://github.com/huggingface/optimum/blob/main/CONTRIBUTING.md
• Most of the work will be extending modeling_ort.py, you can take inspiration from example from

  optimum/optimum/onnxruntime/modeling_ort.py

  Lines 978 to 1020 in 6dc2cf2

  	@add_start_docstrings(
  	"""
  	Onnx Model for image-classification tasks.
  	""",
  	ONNX_MODEL_START_DOCSTRING,
  	)
  	class ORTModelForImageClassification(ORTModel):
  	"""
  	Image Classification model for ONNX.
  	"""
  	# used in from_transformers to export model to onnx
  	export_feature = "image-classification"
  	auto_model_class = AutoModelForImageClassification
  	def __init__(self, model=None, config=None, **kwargs):
  	super().__init__(model, config, **kwargs)
  	# create {name:idx} dict for model outputs
  	self.model_outputs = {output_key.name: idx for idx, output_key in enumerate(self.model.get_outputs())}
  	@add_start_docstrings_to_model_forward(
  	ONNX_IMAGE_INPUTS_DOCSTRING.format("batch_size, num_channels, height, width")
  	+ IMAGE_CLASSIFICATION_EXAMPLE.format(
  	processor_class=_FEATURE_EXTRACTOR_FOR_DOC,
  	model_class="ORTModelForImageClassification",
  	checkpoint="optimum/vit-base-patch16-224",
  	)
  	)
  	def forward(
  	self,
  	pixel_values: torch.Tensor,
  	**kwargs,
  	):
  	# converts pytorch inputs into numpy inputs for onnx
  	onnx_inputs = {
  	"pixel_values": pixel_values.cpu().detach().numpy(),
  	}
  	# run inference
  	outputs = self.model.run(None, onnx_inputs)
  	# converts output to namedtuple for pipelines post-processing
  	return ImageClassifierOutput(
  	logits=torch.from_numpy(outputs[self.model_outputs["logits"]]),
  	)

  . Here one caveat is that models for object detection may have have different classes, for example we have DetrObjectDetectionOutput, OwlViTObjectDetectionOutput, YolosObjectDetectionOutput, etc. As far as I know, there is no map in transformers of a model to its corresponding output class. So a solution could be to output simply a ModelOutput.
• Complete the test_quantization.py, test_modeling.py, test_optimization.py with tests for the new class
• Optionally, add the new class in trainer.py and pipelines.py for pipeline support
• Ping me once you open a PR

Let us know if you are interested in contributing to this!

[shivalikasingh95]

@fxmarty, apologies for the delay in my response. Last week was a bit tight. And sure I'd be happy to contribute to this. The points you have mentioned are very helpful and gives me enough to get started.
However, I'm currently working on another PR for transformers library. It's a model contribution so it's taking some time.
I can take up ORTModelForObjectDetection post finishing that. Hope that's not a problem?

[fxmarty]

Of course no worries! Don't hesitate to open a draft PR once you are on it if you need any help; if a team member wants to implement the class earlier let's follow on in this issue as well 🤗

[shivalikasingh95]

Awesome thanks @fxmarty! I'll keep you posted regarding status on this. Will definitely open up a draft PR once I get started to get some guidance.
And sure, in case any team member wants to take up this task, let me know here :)

[failable]

Any news?