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

[BUG] Remove nested tf.function #391

Closed
arnupretorius opened this issue Feb 4, 2022 · 6 comments · Fixed by #362
Closed

[BUG] Remove nested tf.function #391

arnupretorius opened this issue Feb 4, 2022 · 6 comments · Fixed by #362
Assignees
Labels
bug Something isn't working
Projects

Comments

@arnupretorius
Copy link
Collaborator

arnupretorius commented Feb 4, 2022

Describe the bug
Nested tf.function decorators are causes TF to constantly retrace which is could cause significant performance and memory issues.

Additional context
I think this bug creeped in when we refactored our code to separate forward and backward passes.

Possible Solution
Remove tf.function decorator from the backward pass.

This bug is also related to #77 and #346

@arnupretorius arnupretorius added the bug Something isn't working label Feb 4, 2022
@arnupretorius arnupretorius changed the title [BUG] [BUG] Remove nested tf.function Feb 4, 2022
@arnupretorius arnupretorius self-assigned this Feb 4, 2022
@arnupretorius arnupretorius added this to To do in Mava v0.1.x via automation Feb 4, 2022
@arnupretorius
Copy link
Collaborator Author

arnupretorius commented Feb 4, 2022

Already fixed in #362.

@arnupretorius arnupretorius linked a pull request Feb 4, 2022 that will close this issue
@jcformanek
Copy link
Contributor

Having tf.function over the _policy function in the executor causes retracing as well because _policy is called inside a for loop in the select_actions function. It is better to create a new function called _select_functions and put the for loop inside there and put the tf.function decorator on that function. See MADQN executor in #362

@arnupretorius
Copy link
Collaborator Author

Having tf.function over the _policy function in the executor causes retracing as well because _policy is called inside a for loop in the select_actions function. It is better to create a new function called _select_functions and put the for loop inside there and put the tf.function decorator on that function. See MADQN executor in #362

Thanks @jcformanek! Can you please make an issue for this with the above advice that we look at fixing it for other systems as well.

Mava v0.1.x automation moved this from To do to Done Feb 11, 2022
@jcformanek jcformanek reopened this Feb 11, 2022
Mava v0.1.x automation moved this from Done to In progress Feb 11, 2022
@jcformanek
Copy link
Contributor

Reopening this issue until we fix the problem in all other systems.

@KaleabTessera
Copy link
Contributor

#353 handles this for ppo.

@DriesSmit
Copy link
Contributor

Closing all TF issues as we are depreciating our TF systems.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working
Projects
No open projects
Mava v0.1.x
In progress
Development

Successfully merging a pull request may close this issue.

4 participants