Please allow option to export model to .nn without the action_probs model

[Phong13]

Is your feature request related to a problem? Please describe.
Executing the barracuda models is the biggest performance cost in my project. Using 3 layer by 512 node models, here are the barracuda ModelRunner execution times with different models exported to the .nn file:

• 4.00 ms .nn file with "Action", "Action_probs", "Value_Estimate"
• 2.65 ms .nn file with "Action", Action_probs"
• 2.00 ms .nn file with "Action"

Describe the solution you'd like
When using Barracuda models in inference mode, it is a huge performance savings to be executing leaner .nn files with fewer models in them. Please provide an option to control what gets exported to the .nn files so that it is possible to export the leanest possible .nn file.

Is the "action_probs" model needed in inference mode? Looking at the code in ModelRunner it does not appear that the output from the "action_probs" output is retrieved. Is this model needed by the "action" model somehow? In my testing, the .nn files appear to execute without this model.

Describe alternatives you've considered

Additional context

[Phong13]

Note that it is useful to be able to export the models to separate independent .nn files. In my project I sometimes use the "value_estimate" and sometimes use the "action" output. These are not used at the same time in my project. It is most performant for me to have two models that can be queried separately.

[chriselion]

Hi @Phong13 - sorry for the delay on this.

You're right that it doesn't look like we're using "action_probs" or "value_estimate" at inference time. You could possibly hack the list of output nodes here

ml-agents/ml-agents/mlagents/trainers/tf/model_serialization.py

Lines 43 to 45 in 672b608

	POSSIBLE_OUTPUT_NODES = frozenset(
	["action", "action_probs", "recurrent_out", "value_estimate"]
	)

(although it sounds like that might be what you're already doing). We can look into removing things from this list or making it more configurable in the future.

Would it be possible for you to provide a model file that demonstrates this, so that I can run it by the Barracuda team and make sure there's not a missed optimization opportunity for them? The actual weights don't matter so you can just run training for one step. The easiest format to work with would be the frozen_graph_def.pb file that's generated as part of training. Alternatively, if there's one of our example scenes that's similar to your setup and can be easily modified (either in the training configuration or Agent parameters) that would work too.

Thanks!

[Phong13]

Thanks for looking at this. I did hack the model_serialization.py POSSIBLE_OUTPUT_NODES as you described to produce some .nn files that had different combos of models in them. Here is what they show in the inspector. The value estimate is not normally exported because it doesn't have an identifier in the tensorflow graph (I hacked the trainer to add it back in). Also the value estimate name has changed to "optimizer/value_estimate".

You can see in the inspector that there is more going on in the .nn file graph when there are additional outputs. (44 Layers vs. 36 Layers vs. 27 Layers). I presume this corresponds directly to the effort needed to process these graphs.

If you want to generate some models of your own it is easy to do by hacking the list and in model_serialization.py. The walking biped example scene has similar dimensions to my models (512x3 layers, not as many inputs).

[Phong13]

My dream solution for this would be to configure what outputs are generated. In my case I want two .nn files.

• one would be action only
• one would be value_estimate only

The BehaviourParameters component could have a list of slots for .nn files. I would assign these files to the first and second slots. When I am requesting a decision for an agent I could specifiy which .nn models I want executed. This way I could query for the value_estimate only or the policy output or both at the same time if there was a .nn file configured to include both.

[chriselion]

Hi @Phong13,
I updated the model export so that action_probs will no longer appear as an output. I also added your comments to the internal issue we have tracking your value estimate requests.

[Phong13]

Thanks Chris! :)

Looking forward to the next release.

[chriselion]

Just some bookkeeping - the change to remove action_probs from the model was merged a while ago. Since the discussion on value estimate is tracked elsewhere, I'm going to close this issue.

[github-actions]

This thread has been automatically locked since there has not been any recent activity after it was closed. Please open a new issue for related bugs.