Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Expose more params in main_export for customized ONNX export #2100

Open
wants to merge 5 commits into
base: main
Choose a base branch
from

Conversation

hann-wang
Copy link

@hann-wang hann-wang commented Nov 21, 2024

What does this PR do?

We aim to provide more flexible control over the ONNX exporter, which includes:

  • support for the ROCMExecutionProvider
  • registration of custom operators
  • the ability to use a customized export function besides torch.onnx.export

Currently, there are some challenges in achieving these objectives:

  • providers and session_options are initialized in the fix_dynamic_axes and the _run_validation functions, which offers no interface for registering custom operators or specifying ROCMExecutionProvider.
  • main_export loads models by names or paths, which is unfeasible to export a tailored model.
  • disable_dynamic_axes_fix is not exposed in main_export. If custom operators are not registered in ONNXRuntime, fix_dynamic_axes will fail. To write ONNX models to output, we must be able to disable this function.
  • Many 3rdparty model optimization tools have customized export functions other than torch.onnx.export, making a customizable export function for main_export more desirable.

This PR addresses the aforementioned challenges by:

  • allowing a model object as input to main_export
  • adding disable_dynamic_axes_fix, providers, session_options, and custom_export_fn params to main_export.

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you make sure to update the documentation with your changes?
  • Did you write any new necessary tests?

Who can review?

@fxmarty, @echarlaix, @JingyaHuang, @michaelbenayoun

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant