-
Notifications
You must be signed in to change notification settings - Fork 31.2k
Choose framework for ONNX export #16018
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
Choose framework for ONNX export #16018
Conversation
|
The documentation is not available anymore as the PR was closed or merged. |
lewtun
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks a lot for fixing this @michaelbenayoun - the change looks great! I left some nits about docstrings and a question about running in a pure tensorflow env.
src/transformers/onnx/features.py
Outdated
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What happens here if only tensorflow is installed?
If I understand correctly, since _TASKS_TO_AUTOMODELS is an empty dict by default here, this error won't show the possible values for a pure tensorflow env because we should be accessing _TASKS_TO_TF_AUTOMODELS in that case
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
So I added a method called _validate_framework_choice that should take care of that.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nice!
sgugger
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM!
lewtun
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for iterating on this - LGTM!
src/transformers/onnx/features.py
Outdated
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nice!
c8251ef to
ec48313
Compare
What does this PR do?
This allows to choose which framework to use between PyTorch and TensorFlow for the ONNX export.
Fixes #15990