Implementing loss scaling scheduler callback and schedulers #270
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
Signed-off-by: Lee, Kin Long Kelvin <[email protected]>
Signed-off-by: Kin Long Kelvin Lee <[email protected]>
Signed-off-by: Kin Long Kelvin Lee <[email protected]>
Signed-off-by: Kin Long Kelvin Lee <[email protected]>
Signed-off-by: Kin Long Kelvin Lee <[email protected]>
Signed-off-by: Kin Long Kelvin Lee <[email protected]>
Signed-off-by: Kin Long Kelvin Lee <[email protected]>
This serves as the main interface for controlling the schedules during training
Signed-off-by: Lee, Kin Long Kelvin <[email protected]>
We can set the task scaling values, but if the task key isn't set the data may not be present
Seems like most schedules will need the same configuration
Signed-off-by: Lee, Kin Long Kelvin <[email protected]>
Signed-off-by: Lee, Kin Long Kelvin <[email protected]>
Signed-off-by: Lee, Kin Long Kelvin <[email protected]>
Signed-off-by: Lee, Kin Long Kelvin <[email protected]>
Signed-off-by: Lee, Kin Long Kelvin <[email protected]>
Signed-off-by: Lee, Kin Long Kelvin <[email protected]>
laserkelvin
added
enhancement
|
I think the same tests fail as in #266 |
melo-gonzo
approved these changes
Aug 9, 2024
matsciml/lightning/callbacks.py
Outdated
Comment on lines
1572
to
1585
| target_key = schedule.key | ||
| self._logger.debug( | ||
| f"Attempting to advance {target_key} schedule on step." | ||
| ) | ||
| try: | ||
| new_scaling_value = schedule.step() | ||
| pl_module.task_loss_scaling[target_key] = new_scaling_value | ||
| self._logger.debug( | ||
| f"Advanced {target_key} to new value: {new_scaling_value}" | ||
| ) | ||
| except StopIteration: | ||
| self._logger.warning( | ||
| f"{target_key} has run out of scheduled values; this may be unintentional." | ||
| ) |
There was a problem hiding this comment.
Looks like you could combine this common bit between the on_x_end functions, and pass in a 'step' or 'epoch' string variable to use in the if statement and log message.
There was a problem hiding this comment.
Done in d3b97da
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.
This PR adds a new callback, called
LossScalingScheduler, and scheduler classes that will modify the relative loss weights over the course of training.LinearScalingScheduleandSigmoidScalingSchedule. The former will generate a linear ramp from start to end over steps or epochs, and the latter gives a gradual ramp up in the form of a sigmoid curve.LossScalingScheduleris configured by passing schedules that are mapped to task keys, and for every training step or epoch (set by the schedules), applies a new value of the weighting to the appropriate task.examples/callbacks/loss_scheduling.pyis provided to show how it is configured.This is useful for implementing curricula, where we prioritize learning of different properties over time.