Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feature(models): Add model comm group to predict_step #77

Open
wants to merge 3 commits into
base: develop
Choose a base branch
from
Open

feature(models): Add model comm group to predict_step #77

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e1824cb
added optional model comm group to predict_step
cathalobrien Jan 14, 2025
be9460d
update changelog
cathalobrien Jan 14, 2025
fdd92e6
added **kwargs to predict step
cathalobrien Jan 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions models/CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ Keep it human-readable, your future self will thank you!
- Reduced memory usage when using chunking in the mapper [#84](https://github.com/ecmwf/anemoi-models/pull/84)
- Added `supporting_arrays` argument, which contains arrays to store in checkpoints. [#97](https://github.com/ecmwf/anemoi-models/pull/97)
- Add remappers, e.g. link functions to apply during training to facilitate learning of variables with a difficult distribution [#88](https://github.com/ecmwf/anemoi-models/pull/88)
- 'predict\_step' can take an optional model comm group. [#77](https://github.com/ecmwf/anemoi-core/pull/77)

## [0.4.0](https://github.com/ecmwf/anemoi-models/compare/0.3.0...0.4.0) - Improvements to Model Design

Expand Down
8 changes: 6 additions & 2 deletions models/src/anemoi/models/interface/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,11 @@
# nor does it submit to any jurisdiction.

import uuid
from typing import Optional

import torch
from hydra.utils import instantiate
from torch.distributed.distributed_c10d import ProcessGroup
from torch_geometric.data import HeteroData

from anemoi.models.preprocessing import Processors
Expand Down Expand Up @@ -94,7 +96,9 @@ def _build_model(self) -> None:
# Use the forward method of the model directly
self.forward = self.model.forward

def predict_step(self, batch: torch.Tensor) -> torch.Tensor:
def predict_step(
self, batch: torch.Tensor, model_comm_group: Optional[ProcessGroup] = None, **kwargs
) -> torch.Tensor:
"""Prediction step for the model.

Parameters
Expand All @@ -118,6 +122,6 @@ def predict_step(self, batch: torch.Tensor) -> torch.Tensor:
# batch, timesteps, horizonal space, variables
x = batch[:, 0 : self.multi_step, None, ...] # add dummy ensemble dimension as 3rd index

y_hat = self(x)
y_hat = self(x, model_comm_group)

return self.post_processors(y_hat, in_place=False)
Loading