Expose more params in main_export
for customized ONNX export
#2100
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
What does this PR do?
We aim to provide more flexible control over the ONNX exporter, which includes:
torch.onnx.export
Currently, there are some challenges in achieving these objectives:
providers
andsession_options
are initialized in thefix_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 inmain_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.torch.onnx.export
, making a customizable export function formain_export
more desirable.This PR addresses the aforementioned challenges by:
main_export
disable_dynamic_axes_fix
,providers
,session_options
, andcustom_export_fn
params tomain_export
.Before submitting
Who can review?
@fxmarty, @echarlaix, @JingyaHuang, @michaelbenayoun