Skip to content

Regression & Multitask#452

Open
benedict-armstrong wants to merge 7 commits into
ML4GW:mainfrom
benedict-armstrong:feature/regression_models
Open

Regression & Multitask#452
benedict-armstrong wants to merge 7 commits into
ML4GW:mainfrom
benedict-armstrong:feature/regression_models

Conversation

@benedict-armstrong

Copy link
Copy Markdown

Warning

Work in Progress

This PR builds on top of #451 (most of the changes are from there) and adds two training tasks:

  • regression: projects/train/train/model/regression.py; Designed to be used with waveform_prob=1.0
  • multitask: in projects/train/train/model/multitask.py; Both regression and classification at the same time. So have a regression and a classification head on the same model.

The AframeBase model base class in projects/train/train/model/base.py is refactored to be a general base class for all task types. projects/train/train/model/{classification|regression|multitask}.py inherit from AframeBase and implement specifics for each task.

The two configs in projects/train/configs/regression/ show how to setup the two respective models. Before merging I would remove/refactor these similar to the other training configs.

Benedict Armstrong added 7 commits May 31, 2026 14:25
…kground. Updated datasets, waveform generators and models to accept/return parameters as part of the batch
…lasses. Updated Autoencoder and Supervised models to inherit from AframeClassification. Added new SupervisedMultiTaskAframe and SupervisedRegressionAframe classes for multi-task and regression tasks.
- Introduced RegressionArchitecture and MultiTaskArchitecture classes in regression.py.
- Added RegressionTimeDomainResNet and MultiTaskTimeDomainResNet classes for regression tasks.
- Created multitask.yaml and regression.yaml for multi-task and regression training configurations.
- Enhanced BaseAframeDataset to support parameter transformations.
- Implemented ChirpMass and MassRatio transforms for parameter calculations.
- Updated AframeWandbLogger for improved logging capabilities.
@benedict-armstrong benedict-armstrong changed the title Regression & Multitask models Regression & Multitask Jun 2, 2026

@deepchatterjeeligo deepchatterjeeligo left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Again, I really like breaking up the AframeBase into something simpler. Once you are done with #451, let's rebase this and get this into the regression-datamodules branch

@deepchatterjeeligo

Copy link
Copy Markdown
Contributor

Can this be made against the regression-datamodules branch please?

@kyoon-mit

Copy link
Copy Markdown

Same pull request made to regression-datamodules branch in #503.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants