Feature/22 feature add tendency training #23
Conversation
|
|
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## develop #23 +/- ##
========================================
Coverage 99.82% 99.82%
========================================
Files 21 21
Lines 1158 1158
========================================
Hits 1156 1156
Misses 2 2 ☔ View full report in Codecov by Sentry. |
| @@ -96,8 +97,15 @@ def _build_model(self) -> None: | |||
| config=self.config, data_indices=self.data_indices, graph_data=self.graph_data | |||
| ) | |||
|
|
|||
| # Use the forward method of the model directly | |||
| self.forward = self.model.forward | |||
| def forward(self, x: torch.Tensor, model_comm_group: Optional[ProcessGroup] = None) -> torch.Tensor: | |||
There was a problem hiding this comment.
Wouldn't it make sense to have this better visible, e.g. in the name here? Even if we haven't the alternativ implemented at the moment (but which makes sense to have for anemoi, e.g. with obs the residual is typically not possible).
There was a problem hiding this comment.
I agree, I think instead of self.tendency_mode there should be a self.prediction_mode which can take the values 'state', 'residual' or 'tendency'. I will try to change that when we open the PR with anemoi-training.
|
This should probably be its own interface, rather than replace the working default one, right? |
| assert ( | ||
| len(batch.shape) == 4 | ||
| ), f"The input tensor has an incorrect shape: expected a 4-dimensional tensor, got {batch.shape}!" | ||
| x = self.pre_processors_state(batch[:, 0 : self.multi_step, ...], in_place=False) | ||
|
|
||
| # Dimensions are | ||
| # batch, timesteps, horizonal space, variables | ||
| x = batch[:, 0 : self.multi_step, None, ...] # add dummy ensemble dimension as 3rd index | ||
| # batch, timesteps, horizontal space, variables | ||
| x = x[..., None, :, :] # add dummy ensemble dimension as 3rd index | ||
| if self.prediction_strategy == "tendency": | ||
| tendency_hat = self(x) | ||
| y_hat = self.add_tendency_to_state(x[:, -1, ...], tendency_hat) | ||
| else: | ||
| y_hat = self(x) | ||
| y_hat = self.post_processors_state(y_hat, in_place=False) |
There was a problem hiding this comment.
These multiple if self.prediction_strategy == "tendency" are adding complexity in the main (and currently unique) AnemoiModelInterface class. This will create technical debt and make the code less flexible and more difficult to maintain.
If you want an alternative, this is what you could do:
Since the current AnemoiModelInterface does not have the behaviour you want, you can instead to create another class that behave as you like (copy-paste the whole AnemoiModelInterface in to MyTendencyTrainingAnemoiModelInterface) and use this one instead.
Just after copy-pasting, you realise that you have duplicated code between AnemoiModelInterface and MyTendencyTrainingAnemoiModelInterface. To avoid this, you can put the common code into a mother class BaseAnemoiModelInterface.
There was a problem hiding this comment.
I see that @JesperDramsch had the same opinion on this #23 (comment)
There was a problem hiding this comment.
Sounds good! I’ve already worked on additional interfaces for a different purpose on a local branch, so this should be straightforward to implement. However, if we split the interface, we’ll likely need to divide the current "AnemoiLightningModule" in anemoi-training into two separate modules: one for ForecastingState (ForecastingStateLightningModule) and another for ForecastingTendency (ForecastingTendencyLightningModule).
At the moment, the existing AnemoiLightningModule (which is tied to the current interface) contains logic for handling both state and tendency steps forward.
Closes #22
Describe your changes
< Summary of the changes.>
< Please also include relevant motivation and context. >
Explained in #22
< List any dependencies that are required for this change. >
I had to make PRs/new branches at the following urls:
Type of change
Please delete options that are not relevant.
Checklist before requesting a review
Tag possible reviewers
@clessig