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

Add ONNX Model Export Support #306

Merged
merged 8 commits into from
Nov 15, 2024
Merged

Conversation

nkzawa
Copy link
Contributor

@nkzawa nkzawa commented Oct 21, 2024

This PR introduces support for exporting trained models to the ONNX format in Sample Factory. The feature will be useful in many cases, including for people who want to use their trained models on a different platform or language. It's nice to build the functionality into Sample Factory, since usually, the underlying architecture is something that users don't need to know about.

It's implemented using TorchScript-based ONNX Exporter for now. I tried TorchDynamo as well but couldn’t get it to work.
The usage is described in the added documentation.

@alex-petrenko
Copy link
Owner

Thank you for pushing these!
Looks like there are some merge conflicts and some syntax errors with Python 3.8.
Lmk if you can fix it, otherwise we can probably drop python 3.8 support at this point

@nkzawa
Copy link
Contributor Author

nkzawa commented Oct 23, 2024

Conflict and syntax errors on python3.8 should be fixed. Also, support for the action mask was added in the ONNX export.

EDIT: investigating test errors

@nkzawa
Copy link
Contributor Author

nkzawa commented Oct 26, 2024

@alex-petrenko I think it's ready to get review

Copy link
Owner

@alex-petrenko alex-petrenko left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks amazing! 🔥

@alex-petrenko alex-petrenko merged commit 6810db5 into alex-petrenko:master Nov 15, 2024
4 of 8 checks passed
@nkzawa nkzawa deleted the onnx_export branch November 15, 2024 05:27
fabioseel pushed a commit to fabioseel/sample-factory that referenced this pull request Nov 21, 2024
* support onnx export

* add doc for onnx export

* fix for action_mask

* fix TypeError on python 3.8

* install onnx on ci for mac

* add debug info

* Revert "add debug info"

This reverts commit bc28dbc.

* fix tests
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.

2 participants