Multitask ASE interface #224
Merged
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 <kin.long.kelvin.lee@intel.com>
Signed-off-by: Lee, Kin Long Kelvin <kin.long.kelvin.lee@intel.com>
This isn't actually originally part of the PR scope, but apparently this typo never triggered as an issue until now!
This substantially simplifies the workflow, albeit adds yet another method to multitask modules. The new method simply passes input data into the encoder, and maps it to every single subtask regardless, instead of requiring that the batch shares the dataset keys.
laserkelvin
added
enhancement
|
Ergh, issue with one of the tests... |
This adjusts the logic, albeit maybe inconsistent with the rest of multitask, where we check the incoming batch for dataset names at the top level to determine if it's a multidata batch, instead of relying on the model expectations. This fixes the ase calculate behavior, which would have been mismatched since the module is inherently multidata but the incoming batch is not.
melo-gonzo
reviewed
May 24, 2024
| @@ -2107,7 +2107,7 @@ def __init__( | |||
| if index != 0: | |||
| task.encoder = self.encoder | |||
| # nest the task based on its category | |||
| task_map[dset_name][task.__task__] = task | |||
| task_map[dset_name][task.__class__.__name__] = task | |||
There was a problem hiding this comment.
Was this added for the single data multi-task case?
There was a problem hiding this comment.
Not 100% sure what you mean - I don't think this part of the code matters if it's single data or not?
melo-gonzo
approved these changes
May 24, 2024
Signed-off-by: Lee, Kin Long Kelvin <kin.long.kelvin.lee@intel.com>
|
@melo-gonzo just wanted you to take a look at the example before I merge, and if you wanted to clarify what you meant in your question? |
|
@laserkelvin thanks for the example! I am just curious why you made the change in that line. Feel free to merge ✅ |
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 closes #217 by providing an extended interface for mapping multitask models to
aseworkflows.The main changes include:
matsciml.interfaces.ase.multitask.AbstractStrategyclass, whose children allows us to implement ways to aggregate the outputs of multitasks.AverageTasks, which does a mean aggregation across the different dataset heads to compute energies/forces. Essentially, if you hadForceRegressionTasktrained on multiple datasets, we outputs from all.ase_calculatemethod forMultiTaskModule, which effectively streamlines inference using multitask modules in scenarios like as an ASE calculator. It's named specifically forase, but in the future we can use it as a dedicated inference method.The
AbstractStrategyshould be flexible enough to do other potentially more intelligent aggregations, e.g. a weighted average, MoE approaches, etc.