Added EMA of running weights to the current codebase #853
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
added 3 commits
January 8, 2025 22:35
|
imho the ema stuff should go in the optimizer state, wdyt? |
|
Sounds reasonable to me. This should improve compatibility as well. I will try to implement this. |
|
After trying the implementation, implementing EMA within optimizers seems to limit the generality of this feature. Currently, I have to define specific optimizers like AdamWithEMA while this optimizer is intrinsically the same as Adam. I feel that perhaps it is a better idea to keep the EMA checkpoint in the trainer state because of this. |
|
i see what you're saying, but like, Optax doesn't really have an "adam" right? It's a bunch of transformations and this is just one more? Let me see |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
The syntax works like
'''
--trainer.use_ema True
--trainer.ema_beta 0.995
'''
The current issue is that the current code isn't compatible with the previous version of trainer state.