From 813d36d1bad4919afdd98a4d410bd8ef854cd32c Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Thu, 5 Sep 2024 09:29:58 +0000 Subject: [PATCH 1/6] feat: make model instantiateable --- src/anemoi/models/interface/__init__.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/anemoi/models/interface/__init__.py b/src/anemoi/models/interface/__init__.py index 54c548d..546003a 100644 --- a/src/anemoi/models/interface/__init__.py +++ b/src/anemoi/models/interface/__init__.py @@ -14,7 +14,6 @@ from hydra.utils import instantiate from torch_geometric.data import HeteroData -from anemoi.models.models.encoder_processor_decoder import AnemoiModelEncProcDec from anemoi.models.preprocessing import Processors @@ -73,9 +72,9 @@ def _build_model(self) -> None: self.pre_processors = Processors(processors) self.post_processors = Processors(processors, inverse=True) - # Instantiate the model (Can be generalised to other models in the future, here we use AnemoiModelEncProcDec) - self.model = AnemoiModelEncProcDec( - config=self.config, data_indices=self.data_indices, graph_data=self.graph_data + # Instantiate the model + self.model = instantiate( + self.config.model.model, config=self.config, data_indices=self.data_indices, graph_data=self.graph_data ) # Use the forward method of the model directly From b710a4e6ed897fdd211c5497b5635785ae456889 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Thu, 5 Sep 2024 10:01:57 +0000 Subject: [PATCH 2/6] docs: instantiation explained in changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 73713be..d9737b8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,7 @@ Keep it human-readable, your future self will thank you! - Update CI to inherit from common infrastructue reusable workflows - run downstream-ci only when src and tests folders have changed - New error messages for wrongs graphs. + - Feature: Change model to be instantiatable in the interface, addressing [#28](https://github.com/ecmwf/anemoi-models/issues/28) through [#29](https://github.com/ecmwf/anemoi-models/pulls/29) ### Removed From 7f3ee625778ae84fd4629c52757ae143db003544 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 16 Sep 2024 12:31:26 +0000 Subject: [PATCH 3/6] refactor: rename model config object --- .../models/encoder_processor_decoder.py | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/anemoi/models/models/encoder_processor_decoder.py b/src/anemoi/models/models/encoder_processor_decoder.py index 0f37474..3d3b6b5 100644 --- a/src/anemoi/models/models/encoder_processor_decoder.py +++ b/src/anemoi/models/models/encoder_processor_decoder.py @@ -32,7 +32,7 @@ class AnemoiModelEncProcDec(nn.Module): def __init__( self, *, - config: DotDict, + model_config: DotDict, data_indices: dict, graph_data: HeteroData, ) -> None: @@ -40,8 +40,8 @@ def __init__( Parameters ---------- - config : DotDict - Job configuration + model_config : DotDict + Model configuration data_indices : dict Data indices graph_data : HeteroData @@ -50,15 +50,15 @@ def __init__( super().__init__() self._graph_data = graph_data - self._graph_name_data = config.graph.data - self._graph_name_hidden = config.graph.hidden + self._graph_name_data = model_config.graph.data + self._graph_name_hidden = model_config.graph.hidden self._calculate_shapes_and_indices(data_indices) self._assert_matching_indices(data_indices) - self.multi_step = config.training.multistep_input + self.multi_step = model_config.training.multistep_input - self._define_tensor_sizes(config) + self._define_tensor_sizes(model_config) # Create trainable tensors self._create_trainable_attributes() @@ -67,13 +67,13 @@ def __init__( self._register_latlon("data", self._graph_name_data) self._register_latlon("hidden", self._graph_name_hidden) - self.num_channels = config.model.num_channels + self.num_channels = model_config.model.num_channels input_dim = self.multi_step * self.num_input_channels + self.latlons_data.shape[1] + self.trainable_data_size # Encoder data -> hidden self.encoder = instantiate( - config.model.encoder, + model_config.model.encoder, in_channels_src=input_dim, in_channels_dst=self.latlons_hidden.shape[1] + self.trainable_hidden_size, hidden_dim=self.num_channels, @@ -84,7 +84,7 @@ def __init__( # Processor hidden -> hidden self.processor = instantiate( - config.model.processor, + model_config.model.processor, num_channels=self.num_channels, sub_graph=self._graph_data[(self._graph_name_hidden, "to", self._graph_name_hidden)], src_grid_size=self._hidden_grid_size, @@ -93,7 +93,7 @@ def __init__( # Decoder hidden -> data self.decoder = instantiate( - config.model.decoder, + model_config.model.decoder, in_channels_src=self.num_channels, in_channels_dst=input_dim, hidden_dim=self.num_channels, From 8f0cefb3d52ad3bda99e1d0bab59b5f77007a59c Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 16 Sep 2024 13:18:25 +0000 Subject: [PATCH 4/6] fix: rename config to model_config --- src/anemoi/models/interface/__init__.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/anemoi/models/interface/__init__.py b/src/anemoi/models/interface/__init__.py index 546003a..e6a469c 100644 --- a/src/anemoi/models/interface/__init__.py +++ b/src/anemoi/models/interface/__init__.py @@ -74,7 +74,10 @@ def _build_model(self) -> None: # Instantiate the model self.model = instantiate( - self.config.model.model, config=self.config, data_indices=self.data_indices, graph_data=self.graph_data + self.config.model.model, + model_config=self.config, + data_indices=self.data_indices, + graph_data=self.graph_data, ) # Use the forward method of the model directly From 86a9163770918d2f521aa6f5cac44f25875076a0 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 16 Sep 2024 15:20:18 +0000 Subject: [PATCH 5/6] fix: mark non-recursive --- src/anemoi/models/interface/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/anemoi/models/interface/__init__.py b/src/anemoi/models/interface/__init__.py index 4e498f9..aba62a2 100644 --- a/src/anemoi/models/interface/__init__.py +++ b/src/anemoi/models/interface/__init__.py @@ -78,6 +78,7 @@ def _build_model(self) -> None: model_config=self.config, data_indices=self.data_indices, graph_data=self.graph_data, + _recursive_=False, # Disables recursive instantiation by Hydra ) # Use the forward method of the model directly From 32577ff672219496cc3296ed62eea1e2c17c662f Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 16 Sep 2024 15:24:17 +0000 Subject: [PATCH 6/6] docs: changelog --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 84c3f1c..6a9c8f9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,7 +22,7 @@ Keep it human-readable, your future self will thank you! - Update CI to inherit from common infrastructue reusable workflows - run downstream-ci only when src and tests folders have changed - New error messages for wrongs graphs. -- Feature: Change model to be instantiatable in the interface, addressing [#28](https://github.com/ecmwf/anemoi-models/issues/28) through [#29](https://github.com/ecmwf/anemoi-models/pulls/29) +- Feature: Change model to be instantiatable in the interface, addressing [#28](https://github.com/ecmwf/anemoi-models/issues/28) through [#45](https://github.com/ecmwf/anemoi-models/pulls/45) ### Removed