-
Notifications
You must be signed in to change notification settings - Fork 611
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
Make MultiOptimizer serializable #2719
Make MultiOptimizer serializable #2719
Conversation
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
You are owner of some files modified in this pull request. |
The current implementation of MultiOptimizer puts grads and vars inside the optimizer_specs variable. Since the grads and vars are not serializable, any model that uses MultiOptimizer class cannot be saved or checkpointed. This PR will only extract the necessary optimizer_specs that are useful for re-instantiating the MultiOptimizer instance in get_config() method, leading to a serializable MultiOptimizer implementation.
2529b5c
to
ba536e0
Compare
Can you add a serialization test? |
The serialization test has already been there (code ref). The problem is that the grads and vars are not there until the model has been trained for at least one iteration. Any suggestion on how to deal with this? |
You could extend the test. |
Yes you are right. Sorry I didn't read enough code to realize we can fit the model during testing. I have added a test that check the serializability. And I have verified that without this PR the test won't pass. |
@bhack Is there any more thing I can do to help facilitate merging of this PR? |
Thanks |
Description
The current implementation of
MultiOptimizer
puts grads and vars inside theoptimizer_specs
variable. Since the grads and vars are not serializable, any model that usesMultiOptimizer
class cannot be saved or checkpointed. This PR will only extract the necessaryoptimizer_specs
components that are useful for re-instantiating theMultiOptimizer
instance insideget_config()
method, leading to a serializableMultiOptimizer
implementation.Fixes # (issue)
Type of change
Checklist:
How Has This Been Tested?
If you're adding a bugfix or new feature please describe the tests that you ran to verify your changes:
*
I have tested
model.save()
method after this change. Prior to this change, a trained model (has run training iteration for at least 1 step) withMultiOptimizer
cannot be saved because of the un-serializable grads and vars. With the fix, I can verify thatmodel.save()
method works.