Support for ONNX export #3101
Support for ONNX export #3101
Conversation
UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/BarracudaModelParamLoader.cs
Outdated
Show resolved
Hide resolved
|
I cancelled the yamato tests. Will add changelog entries after the release branch is merged to master. |
ml-agents/mlagents/serialization.py
Outdated
| @@ -0,0 +1,205 @@ | |||
| import logging | |||
There was a problem hiding this comment.
would it make sense to call this file onnx_serialization.py?
There was a problem hiding this comment.
It also handles the tensorflow_to_barracuda conversion (those methods were migrated out of tf_policy.py). Maybemodel_serialization.py instead?
There was a problem hiding this comment.
eh, guess not. looks like a lot more is going on here.
There was a problem hiding this comment.
i like model_serialization.py
|
This doesn't yet do ONNX -> Barracuda right? The export part LGTM |
| self.get_policy(name_behavior_id).export_model() | ||
| policy = self.get_policy(name_behavior_id) | ||
| settings = SerializationSettings(policy.model_path, policy.brain.brain_name) | ||
| export_policy_model(settings, policy.graph, policy.sess) |
There was a problem hiding this comment.
👍 Do like moving exporting out of Policy
|
@ervteng Barracuda importer will handle everything, we just need to export the ONNX. |
| @@ -27,10 +27,17 @@ but we only tested for the following platforms : | |||
| * iOS | |||
| * Android | |||
|
|
|||
| ## Supported formats | |||
| There are currently two supported model formats: | |||
Adds support for exporting to ONNX format. This will be disabled by default, and enabled if the
tf2onnxpackage is installed. A recent version (>=1.5.5) is recommended. Note that tf2onnx is not currently compatible with tensorflow 2.0.0 or higher.ONNX files generated from ML-Agents can only be imported with Barracuda v0.6.0 or later.
Other notable changes: