From dbee83b3b9b48a920f0f6b23743d186213fd0b31 Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Wed, 11 Sep 2024 17:24:04 +0200 Subject: [PATCH 1/6] Docs/changelog update 0.3.0 (#40) * [create-pull-request] automated change * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update CHANGELOG.md --------- Co-authored-by: JesperDramsch <2620316+JesperDramsch@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- CHANGELOG.md | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5433ccc..55fba06 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,21 +8,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 Please add your functional changes to the appropriate section in the PR. Keep it human-readable, your future self will thank you! -## [Unreleased] +## [Unreleased](https://github.com/ecmwf/anemoi-models/compare/0.3.0...HEAD) + +## [0.3.0](https://github.com/ecmwf/anemoi-models/compare/0.2.1...0.3.0) - Remapping of (meteorological) Variables ### Added - - CI workflow to update the changelog on release - - Remapper: Preprocessor for remapping one variable to multiple ones. Includes changes to the data indices since the remapper changes the number of variables. With optional config keywords. + +- CI workflow to update the changelog on release +- Remapper: Preprocessor for remapping one variable to multiple ones. Includes changes to the data indices since the remapper changes the number of variables. With optional config keywords. ### Changed - - 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. +- 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. ### Removed -## [0.2.1] - Dependency update +## [0.2.1](https://github.com/ecmwf/anemoi-models/compare/0.2.0...0.2.1) - Dependency update ### Added @@ -33,7 +36,7 @@ Keep it human-readable, your future self will thank you! - anemoi-datasets dependency -## [0.2.0] - Support Heterodata +## [0.2.0](https://github.com/ecmwf/anemoi-models/compare/0.1.0...0.2.0) - Support Heterodata ### Added @@ -43,15 +46,12 @@ Keep it human-readable, your future self will thank you! - Updated to support new PyTorch Geometric HeteroData structure (defined by `anemoi-graphs` package). -## [0.1.0] - Initial Release +## [0.1.0](https://github.com/ecmwf/anemoi-models/releases/tag/0.1.0) - Initial Release ### Added + - Documentation - Initial code release with models, layers, distributed, preprocessing, and data_indices - Added Changelog -[unreleased]: https://github.com/ecmwf/anemoi-models/compare/0.2.1...HEAD -[0.2.1]: https://github.com/ecmwf/anemoi-models/compare/0.2.0...0.2.1 -[0.2.0]: https://github.com/ecmwf/anemoi-models/compare/0.1.0...0.2.0 -[0.1.0]: https://github.com/ecmwf/anemoi-models/releases/tag/0.1.0 From 43846f042a44699772b4f6d9830bd3f85b2ffd4c Mon Sep 17 00:00:00 2001 From: Helen Theissen Date: Wed, 18 Sep 2024 11:14:40 +0100 Subject: [PATCH 2/6] Chore/multiple fixes ci precommit (#41) * fix: change pre-cmmit autoupdate schedule to monthly * fix: change the merge strategy for Changelog to Union * fix: add .envrc to .gitignore * ci: ignore pre-commit-config and readthedocs for changelog updates * ci: fix to correct hpc workflow call * fix: update precommit config * chore: update pre-commits * feat: add codeowners file * chore: update dependencies * ci: add hpc-config * docs: changelog * fix: respond to review comments --------- Co-authored-by: Jesper Dramsch --- .gitattributes | 1 + .github/CODEOWNERS | 6 ++++ .github/ci-hpc-config.yml | 7 +++++ .github/workflows/changelog-pr-update.yml | 3 ++ .github/workflows/ci.yml | 2 +- .pre-commit-config.yaml | 31 +++++++++++++++------ CHANGELOG.md | 5 ++++ pyproject.toml | 34 ++++------------------- 8 files changed, 52 insertions(+), 37 deletions(-) create mode 100644 .gitattributes create mode 100644 .github/CODEOWNERS create mode 100644 .github/ci-hpc-config.yml diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..a19ade0 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +CHANGELOG.md merge=union diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 0000000..a2c619f --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1,6 @@ +# CODEOWNERS file + +# Protect workflow files +/.github/ @theissenhelen @jesperdramsch @gmertes +/.pre-commit-config.yaml @theissenhelen @jesperdramsch @gmertes +/pyproject.toml @theissenhelen @jesperdramsch @gmertes diff --git a/.github/ci-hpc-config.yml b/.github/ci-hpc-config.yml new file mode 100644 index 0000000..5c27b03 --- /dev/null +++ b/.github/ci-hpc-config.yml @@ -0,0 +1,7 @@ +build: + python: '3.10' + modules: + - ninja + python_dependencies: + - ecmwf/anemoi-utils@develop + parallel: 64 diff --git a/.github/workflows/changelog-pr-update.yml b/.github/workflows/changelog-pr-update.yml index 43acb1c..e7ed9a2 100644 --- a/.github/workflows/changelog-pr-update.yml +++ b/.github/workflows/changelog-pr-update.yml @@ -2,6 +2,9 @@ name: Check Changelog Update on PR on: pull_request: types: [assigned, opened, synchronize, reopened, labeled, unlabeled] + paths-ignore: + - .pre-commit-config.yaml + - .readthedocs.yaml jobs: Check-Changelog: name: Check Changelog Action diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5867ee0..8b2926b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -37,7 +37,7 @@ jobs: downstream-ci-hpc: name: downstream-ci-hpc if: ${{ !github.event.pull_request.head.repo.fork && github.event.action != 'labeled' || github.event.label.name == 'approved-for-ci' }} - uses: ecmwf-actions/downstream-ci/.github/workflows/downstream-ci.yml@main + uses: ecmwf-actions/downstream-ci/.github/workflows/downstream-ci-hpc.yml@main with: anemoi-models: ecmwf/anemoi-models@${{ github.event.pull_request.head.sha || github.sha }} secrets: inherit diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c042b1f..f3c3962 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -20,6 +20,12 @@ repos: - id: no-commit-to-branch # Prevent committing to main / master - id: check-added-large-files # Check for large files added to git - id: check-merge-conflict # Check for files that contain merge conflict +- repo: https://github.com/pre-commit/pygrep-hooks + rev: v1.10.0 # Use the ref you want to point at + hooks: + - id: python-use-type-annotations # Check for missing type annotations + - id: python-check-blanket-noqa # Check for # noqa: all + - id: python-no-log-warn # Check for log.warn - repo: https://github.com/psf/black-pre-commit-mirror rev: 24.8.0 hooks: @@ -34,7 +40,7 @@ repos: - --force-single-line-imports - --profile black - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.6.3 + rev: v0.6.4 hooks: - id: ruff # Next line if for documenation cod snippets @@ -45,7 +51,7 @@ repos: - --exit-non-zero-on-fix - --preview - repo: https://github.com/sphinx-contrib/sphinx-lint - rev: v0.9.1 + rev: v1.0.0 hooks: - id: sphinx-lint # For now, we use it. But it does not support a lot of sphinx features @@ -59,12 +65,21 @@ repos: hooks: - id: docconvert args: ["numpy"] -- repo: https://github.com/b8raoult/optional-dependencies-all - rev: "0.0.6" - hooks: - - id: optional-dependencies-all - args: ["--inplace", "--exclude-keys=dev,docs,tests", "--group=dev=all,docs,tests"] - repo: https://github.com/tox-dev/pyproject-fmt - rev: "2.2.1" + rev: "2.2.3" hooks: - id: pyproject-fmt +- repo: https://github.com/jshwi/docsig # Check docstrings against function sig + rev: v0.60.1 + hooks: + - id: docsig + args: + - --ignore-no-params # Allow docstrings without parameters + - --check-dunders # Check dunder methods + - --check-overridden # Check overridden methods + - --check-protected # Check protected methods + - --check-class # Check class docstrings + - --disable=E113 # Disable empty docstrings + - --summary # Print a summary +ci: + autoupdate_schedule: monthly diff --git a/CHANGELOG.md b/CHANGELOG.md index 55fba06..e6553c7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,12 +16,17 @@ Keep it human-readable, your future self will thank you! - CI workflow to update the changelog on release - Remapper: Preprocessor for remapping one variable to multiple ones. Includes changes to the data indices since the remapper changes the number of variables. With optional config keywords. +- Codeowners file +- Pygrep precommit hooks +- Docsig precommit hooks +- Changelog merge strategy ### Changed - 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. +- Bugfixes for CI ### Removed diff --git a/pyproject.toml b/pyproject.toml index 05a99c4..214f82c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,20 +11,13 @@ [build-system] build-backend = "setuptools.build_meta" -requires = [ - "setuptools>=61", - "setuptools-scm>=8", -] +requires = [ "setuptools>=61", "setuptools-scm>=8" ] [project] name = "anemoi-models" description = "A package to hold various functions to support training of ML models." readme = "README.md" -keywords = [ - "ai", - "models", - "tools", -] +keywords = [ "ai", "models", "tools" ] license = { file = "LICENSE" } authors = [ @@ -47,9 +40,7 @@ classifiers = [ "Programming Language :: Python :: Implementation :: PyPy", ] -dynamic = [ - "version", -] +dynamic = [ "version" ] dependencies = [ "anemoi-utils>=0.1.9", "einops>=0.6.1", @@ -57,19 +48,9 @@ dependencies = [ "torch>=2.2", "torch-geometric>=2.3,<2.5", ] -optional-dependencies.all = [ -] +optional-dependencies.all = [ ] -optional-dependencies.dev = [ - "hypothesis", - "nbsphinx", - "pandoc", - "pytest", - "rstfmt", - "sphinx", - "sphinx-argparse<0.5", - "sphinx-rtd-theme", -] +optional-dependencies.dev = [ "anemoi-models[all,docs,tests]" ] optional-dependencies.docs = [ "nbsphinx", @@ -80,10 +61,7 @@ optional-dependencies.docs = [ "sphinx-rtd-theme", ] -optional-dependencies.tests = [ - "hypothesis", - "pytest", -] +optional-dependencies.tests = [ "hypothesis", "pytest" ] urls.Documentation = "https://anemoi-models.readthedocs.io/" urls.Homepage = "https://github.com/ecmwf/anemoi-models/" From 02192661aede10baef6b03682058c01378bdac29 Mon Sep 17 00:00:00 2001 From: Helen Theissen Date: Wed, 18 Sep 2024 18:37:06 +0100 Subject: [PATCH 3/6] 11 add configurability to dropout in multiheadselfattention module (#12) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: add configurability to dropout in MultiHeadSelfAttention Co-authored-by: Rilwan (Akanni) Adewoyin <18564167+Rilwan-Adewoyin@users.noreply.github.com> * test: adjust to dropout_p * doc: update changelog * Feature/integrate reusable workflows (#16) * ci: add public pr label * ci: add readthedocs update check * ci: add downstream ci * ci: add ci-config * chore(deps): remove unused dependency * docs: update changelog * ci: switch to main * chore: changelog 0.2.1 * Update error messages from invalid sub_graph in model instantiation (#20) * ci: inherit pypi publish flow (#17) * ci: inherit pypi publish flow Co-authored-by: Helen Theissen * docs: add to changelog * fix: typo in reusable workflow * fix: another typo * chore: bump actions/setup-python to v5 * ci: run downstream-ci for changes in src and tests * docs: update changelog --------- Co-authored-by: Helen Theissen * Update CHANGELOG.md to KeepChangelog format * [pre-commit.ci] pre-commit autoupdate (#25) updates: - [github.com/psf/black-pre-commit-mirror: 24.4.2 → 24.8.0](https://github.com/psf/black-pre-commit-mirror/compare/24.4.2...24.8.0) - [github.com/astral-sh/ruff-pre-commit: v0.4.6 → v0.6.2](https://github.com/astral-sh/ruff-pre-commit/compare/v0.4.6...v0.6.2) - [github.com/tox-dev/pyproject-fmt: 2.1.3 → 2.2.1](https://github.com/tox-dev/pyproject-fmt/compare/2.1.3...2.2.1) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * Ci/changelog-release-updater (#26) * ci: add changelof release updater * docs: update changelog * Feature/integrate reusable workflows (#16) * ci: add public pr label * ci: add readthedocs update check * ci: add downstream ci * ci: add ci-config * chore(deps): remove unused dependency * docs: update changelog * ci: switch to main * chore: changelog 0.2.1 * Update error messages from invalid sub_graph in model instantiation (#20) * ci: inherit pypi publish flow (#17) * ci: inherit pypi publish flow Co-authored-by: Helen Theissen * docs: add to changelog * fix: typo in reusable workflow * fix: another typo * chore: bump actions/setup-python to v5 * ci: run downstream-ci for changes in src and tests * docs: update changelog --------- Co-authored-by: Helen Theissen * Update CHANGELOG.md to KeepChangelog format * Ci/changelog-release-updater (#26) * ci: add changelof release updater * docs: update changelog --------- Co-authored-by: Rilwan (Akanni) Adewoyin <18564167+Rilwan-Adewoyin@users.noreply.github.com> Co-authored-by: Gert Mertes Co-authored-by: Mario Santa Cruz <48736305+JPXKQX@users.noreply.github.com> Co-authored-by: Jesper Dramsch Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- CHANGELOG.md | 3 +++ src/anemoi/models/layers/attention.py | 15 +++++++++++---- src/anemoi/models/layers/block.py | 12 ++++++++++-- src/anemoi/models/layers/chunk.py | 4 ++++ src/anemoi/models/layers/processor.py | 4 ++++ tests/layers/block/test_block_transformer.py | 13 ++++++++++--- tests/layers/chunk/test_chunk_transformer.py | 4 ++++ .../processor/test_transformer_processor.py | 6 ++++++ tests/layers/test_attention.py | 16 ++++++++++------ 9 files changed, 62 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e6553c7..ec4db72 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,8 @@ Keep it human-readable, your future self will thank you! ### Added +- CI workflow to update the changelog on release +- configurabilty of the dropout probability in the the MultiHeadSelfAttention module - CI workflow to update the changelog on release - Remapper: Preprocessor for remapping one variable to multiple ones. Includes changes to the data indices since the remapper changes the number of variables. With optional config keywords. - Codeowners file @@ -21,6 +23,7 @@ Keep it human-readable, your future self will thank you! - Docsig precommit hooks - Changelog merge strategy + ### Changed - Update CI to inherit from common infrastructue reusable workflows diff --git a/src/anemoi/models/layers/attention.py b/src/anemoi/models/layers/attention.py index 2063ad0..931e098 100644 --- a/src/anemoi/models/layers/attention.py +++ b/src/anemoi/models/layers/attention.py @@ -40,7 +40,7 @@ def __init__( bias: bool = False, is_causal: bool = False, window_size: Optional[int] = None, - dropout: float = 0.0, + dropout_p: float = 0.0, ): super().__init__() @@ -48,11 +48,11 @@ def __init__( embed_dim % num_heads == 0 ), f"Embedding dimension ({embed_dim}) must be divisible by number of heads ({num_heads})" - self.dropout = dropout self.num_heads = num_heads self.embed_dim = embed_dim self.head_dim = embed_dim // num_heads # q k v self.window_size = (window_size, window_size) # flash attention + self.dropout_p = dropout_p self.is_causal = is_causal self.lin_qkv = nn.Linear(embed_dim, 3 * embed_dim, bias=bias) @@ -86,15 +86,22 @@ def forward( query = shard_heads(query, shapes=shapes, mgroup=model_comm_group) key = shard_heads(key, shapes=shapes, mgroup=model_comm_group) value = shard_heads(value, shapes=shapes, mgroup=model_comm_group) + dropout_p = self.dropout_p if self.training else 0.0 if _FLASH_ATTENTION_AVAILABLE: query, key, value = ( einops.rearrange(t, "batch heads grid vars -> batch grid heads vars") for t in (query, key, value) ) - out = self.attention(query, key, value, causal=False, window_size=self.window_size) + out = self.attention(query, key, value, causal=False, window_size=self.window_size, dropout_p=dropout_p) out = einops.rearrange(out, "batch grid heads vars -> batch heads grid vars") else: - out = self.attention(query, key, value, is_causal=False) # expects (batch heads grid variable) format + out = self.attention( + query, + key, + value, + is_causal=False, + dropout_p=dropout_p, + ) # expects (batch heads grid variable) format out = shard_sequence(out, shapes=shapes, mgroup=model_comm_group) out = einops.rearrange(out, "batch heads grid vars -> (batch grid) (heads vars)") diff --git a/src/anemoi/models/layers/block.py b/src/anemoi/models/layers/block.py index ba29607..7fd3627 100644 --- a/src/anemoi/models/layers/block.py +++ b/src/anemoi/models/layers/block.py @@ -55,7 +55,15 @@ def forward( class TransformerProcessorBlock(BaseBlock): """Transformer block with MultiHeadSelfAttention and MLPs.""" - def __init__(self, num_channels, hidden_dim, num_heads, activation, window_size: int): + def __init__( + self, + num_channels: int, + hidden_dim: int, + num_heads: int, + activation: str, + window_size: int, + dropout_p: float = 0.0, + ): super().__init__() try: @@ -72,7 +80,7 @@ def __init__(self, num_channels, hidden_dim, num_heads, activation, window_size: window_size=window_size, bias=False, is_causal=False, - dropout=0.0, + dropout_p=dropout_p, ) self.mlp = nn.Sequential( diff --git a/src/anemoi/models/layers/chunk.py b/src/anemoi/models/layers/chunk.py index 61dec34..87d0ac7 100644 --- a/src/anemoi/models/layers/chunk.py +++ b/src/anemoi/models/layers/chunk.py @@ -73,6 +73,7 @@ def __init__( num_heads: int = 16, mlp_hidden_ratio: int = 4, activation: str = "GELU", + dropout_p: float = 0.0, ) -> None: """Initialize TransformerProcessor. @@ -88,6 +89,8 @@ def __init__( ratio of mlp hidden dimension to embedding dimension, default 4 activation : str, optional Activation function, by default "GELU" + dropout_p: float + Dropout probability used for multi-head self attention, default 0.0 """ super().__init__(num_channels=num_channels, num_layers=num_layers) @@ -98,6 +101,7 @@ def __init__( num_heads=num_heads, activation=activation, window_size=window_size, + dropout_p=dropout_p, ) def forward( diff --git a/src/anemoi/models/layers/processor.py b/src/anemoi/models/layers/processor.py index bb33609..6ba8eb1 100644 --- a/src/anemoi/models/layers/processor.py +++ b/src/anemoi/models/layers/processor.py @@ -95,6 +95,7 @@ def __init__( cpu_offload: bool = False, num_heads: int = 16, mlp_hidden_ratio: int = 4, + dropout_p: float = 0.1, **kwargs, ) -> None: """Initialize TransformerProcessor. @@ -113,6 +114,8 @@ def __init__( ratio of mlp hidden dimension to embedding dimension, default 4 activation : str, optional Activation function, by default "GELU" + dropout_p: float, optional + Dropout probability used for multi-head self attention, default 0.0 """ super().__init__( num_channels=num_channels, @@ -133,6 +136,7 @@ def __init__( num_layers=self.chunk_size, window_size=window_size, activation=activation, + dropout_p=dropout_p, ) self.offload_layers(cpu_offload) diff --git a/tests/layers/block/test_block_transformer.py b/tests/layers/block/test_block_transformer.py index 97c274f..2e63386 100644 --- a/tests/layers/block/test_block_transformer.py +++ b/tests/layers/block/test_block_transformer.py @@ -29,11 +29,14 @@ class TestTransformerProcessorBlock: num_heads=st.integers(min_value=1, max_value=10), activation=st.sampled_from(["ReLU", "GELU", "Tanh"]), window_size=st.integers(min_value=1, max_value=512), + dropout_p=st.floats(min_value=0.0, max_value=1.0), ) @settings(max_examples=10) - def test_init(self, factor_attention_heads, hidden_dim, num_heads, activation, window_size): + def test_init(self, factor_attention_heads, hidden_dim, num_heads, activation, window_size, dropout_p): num_channels = num_heads * factor_attention_heads - block = TransformerProcessorBlock(num_channels, hidden_dim, num_heads, activation, window_size) + block = TransformerProcessorBlock( + num_channels, hidden_dim, num_heads, activation, window_size, dropout_p=dropout_p + ) assert isinstance(block, TransformerProcessorBlock) assert isinstance(block.layer_norm1, nn.LayerNorm) @@ -49,6 +52,7 @@ def test_init(self, factor_attention_heads, hidden_dim, num_heads, activation, w window_size=st.integers(min_value=1, max_value=512), shapes=st.lists(st.integers(min_value=1, max_value=10), min_size=3, max_size=3), batch_size=st.integers(min_value=1, max_value=40), + dropout_p=st.floats(min_value=0.0, max_value=1.0), ) @settings(max_examples=10) def test_forward_output( @@ -60,9 +64,12 @@ def test_forward_output( window_size, shapes, batch_size, + dropout_p, ): num_channels = num_heads * factor_attention_heads - block = TransformerProcessorBlock(num_channels, hidden_dim, num_heads, activation, window_size) + block = TransformerProcessorBlock( + num_channels, hidden_dim, num_heads, activation, window_size, dropout_p=dropout_p + ) x = torch.randn((batch_size, num_channels)) diff --git a/tests/layers/chunk/test_chunk_transformer.py b/tests/layers/chunk/test_chunk_transformer.py index 1fe7c6d..5449e97 100644 --- a/tests/layers/chunk/test_chunk_transformer.py +++ b/tests/layers/chunk/test_chunk_transformer.py @@ -20,6 +20,7 @@ def init(self): mlp_hidden_ratio: int = 4 activation: str = "GELU" window_size: int = 13 + dropout_p: float = 0.1 # num_heads must be evenly divisible by num_channels for MHSA return ( @@ -29,6 +30,7 @@ def init(self): mlp_hidden_ratio, activation, window_size, + dropout_p, ) @pytest.fixture @@ -40,6 +42,7 @@ def processor_chunk(self, init): mlp_hidden_ratio, activation, window_size, + dropout_p, ) = init return TransformerProcessorChunk( num_channels=num_channels, @@ -48,6 +51,7 @@ def processor_chunk(self, init): mlp_hidden_ratio=mlp_hidden_ratio, activation=activation, window_size=window_size, + dropout_p=dropout_p, ) def test_all_blocks(self, processor_chunk): diff --git a/tests/layers/processor/test_transformer_processor.py b/tests/layers/processor/test_transformer_processor.py index d359c27..305af41 100644 --- a/tests/layers/processor/test_transformer_processor.py +++ b/tests/layers/processor/test_transformer_processor.py @@ -21,6 +21,7 @@ def transformer_processor_init(): cpu_offload = False num_heads = 16 mlp_hidden_ratio = 4 + dropout_p = 0.1 return ( num_layers, window_size, @@ -30,6 +31,7 @@ def transformer_processor_init(): cpu_offload, num_heads, mlp_hidden_ratio, + dropout_p, ) @@ -44,6 +46,7 @@ def transformer_processor(transformer_processor_init): cpu_offload, num_heads, mlp_hidden_ratio, + dropout_p, ) = transformer_processor_init return TransformerProcessor( num_layers=num_layers, @@ -54,6 +57,7 @@ def transformer_processor(transformer_processor_init): cpu_offload=cpu_offload, num_heads=num_heads, mlp_hidden_ratio=mlp_hidden_ratio, + dropout_p=dropout_p, ) @@ -67,6 +71,7 @@ def test_transformer_processor_init(transformer_processor, transformer_processor _cpu_offload, _num_heads, _mlp_hidden_ratio, + _dropout_p, ) = transformer_processor_init assert isinstance(transformer_processor, TransformerProcessor) assert transformer_processor.num_chunks == num_chunks @@ -84,6 +89,7 @@ def test_transformer_processor_forward(transformer_processor, transformer_proces _cpu_offload, _num_heads, _mlp_hidden_ratio, + _dropout_p, ) = transformer_processor_init gridsize = 100 batch_size = 1 diff --git a/tests/layers/test_attention.py b/tests/layers/test_attention.py index ffeaebc..9457317 100644 --- a/tests/layers/test_attention.py +++ b/tests/layers/test_attention.py @@ -18,17 +18,19 @@ @given( num_heads=st.integers(min_value=1, max_value=50), embed_dim_multiplier=st.integers(min_value=1, max_value=10), + dropout_p=st.floats(min_value=0.0, max_value=1.0), ) -def test_multi_head_self_attention_init(num_heads, embed_dim_multiplier): +def test_multi_head_self_attention_init(num_heads, embed_dim_multiplier, dropout_p): embed_dim = ( num_heads * embed_dim_multiplier ) # TODO: Make assert in MHSA to check if embed_dim is divisible by num_heads - mhsa = MultiHeadSelfAttention(num_heads, embed_dim) + mhsa = MultiHeadSelfAttention(num_heads, embed_dim, dropout_p=dropout_p) assert isinstance(mhsa, nn.Module) assert mhsa.num_heads == num_heads assert mhsa.embed_dim == embed_dim assert mhsa.head_dim == embed_dim // num_heads + assert dropout_p == mhsa.dropout_p @pytest.mark.gpu @@ -36,11 +38,12 @@ def test_multi_head_self_attention_init(num_heads, embed_dim_multiplier): batch_size=st.integers(min_value=1, max_value=64), num_heads=st.integers(min_value=1, max_value=20), embed_dim_multiplier=st.integers(min_value=1, max_value=10), + dropout_p=st.floats(min_value=0.0, max_value=1.0), ) @settings(deadline=None) -def test_multi_head_self_attention_forward(batch_size, num_heads, embed_dim_multiplier): +def test_multi_head_self_attention_forward(batch_size, num_heads, embed_dim_multiplier, dropout_p): embed_dim = num_heads * embed_dim_multiplier - mhsa = MultiHeadSelfAttention(num_heads, embed_dim) + mhsa = MultiHeadSelfAttention(num_heads, embed_dim, dropout_p=dropout_p) x = torch.randn(batch_size * 2, embed_dim) shapes = [list(x.shape)] @@ -54,10 +57,11 @@ def test_multi_head_self_attention_forward(batch_size, num_heads, embed_dim_mult batch_size=st.integers(min_value=1, max_value=64), num_heads=st.integers(min_value=1, max_value=20), embed_dim_multiplier=st.integers(min_value=1, max_value=10), + dropout_p=st.floats(min_value=0.0, max_value=1.0), ) -def test_multi_head_self_attention_backward(batch_size, num_heads, embed_dim_multiplier): +def test_multi_head_self_attention_backward(batch_size, num_heads, embed_dim_multiplier, dropout_p): embed_dim = num_heads * embed_dim_multiplier - mhsa = MultiHeadSelfAttention(num_heads, embed_dim) + mhsa = MultiHeadSelfAttention(num_heads, embed_dim, dropout_p=dropout_p) x = torch.randn(batch_size * 2, embed_dim, requires_grad=True) shapes = [list(x.shape)] From 90ef59c2abb41b74a273b6b8863a64bdb26fc7b6 Mon Sep 17 00:00:00 2001 From: gabrieloks <116646686+gabrieloks@users.noreply.github.com> Date: Mon, 23 Sep 2024 13:55:09 +0100 Subject: [PATCH 4/6] activation functions added for bounded outputs (#14) * activation functions added for bounded outputs * generalised fraction normalisation * precommit changes * precommit changes * chore: mv bounding to post-processors * feat: make bounding strategies torch Module * refactor: pre-build indices during init * refactor: nn.module in place modification * refactor: mv bounding to layers * refactor: implement bounding ModuleList * docs: add changelog * refactor: mv bounding config to models * feat: enable 1-1 variable remapping in preprocessors * test: create tests and fix code * test: add a test for hydra instantiating of bounding * fix: naming * refactor: reduce verboseness * docs: add comments * fix: inject name_to_index on initiation * fixed reading of statistics remapping * revert * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update test_preprocessor_normalizer.py * Update encoder_processor_decoder.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * docs: added a comment on special keys * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * docs: add docstrings * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: changelog --------- Co-authored-by: Jesper Dramsch --- CHANGELOG.md | 21 ++-- src/anemoi/models/layers/bounding.py | 115 ++++++++++++++++++ .../models/encoder_processor_decoder.py | 15 +++ src/anemoi/models/preprocessing/__init__.py | 20 ++- src/anemoi/models/preprocessing/normalizer.py | 22 ++++ tests/layers/test_bounding.py | 92 ++++++++++++++ .../test_preprocessor_normalizer.py | 43 +++++++ 7 files changed, 319 insertions(+), 9 deletions(-) create mode 100644 src/anemoi/models/layers/bounding.py create mode 100644 tests/layers/test_bounding.py diff --git a/CHANGELOG.md b/CHANGELOG.md index ec4db72..7811ef3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,26 +10,31 @@ Keep it human-readable, your future self will thank you! ## [Unreleased](https://github.com/ecmwf/anemoi-models/compare/0.3.0...HEAD) -## [0.3.0](https://github.com/ecmwf/anemoi-models/compare/0.2.1...0.3.0) - Remapping of (meteorological) Variables - ### Added - -- CI workflow to update the changelog on release -- configurabilty of the dropout probability in the the MultiHeadSelfAttention module -- CI workflow to update the changelog on release -- Remapper: Preprocessor for remapping one variable to multiple ones. Includes changes to the data indices since the remapper changes the number of variables. With optional config keywords. - Codeowners file - Pygrep precommit hooks - Docsig precommit hooks - Changelog merge strategy +- configurabilty of the dropout probability in the the MultiHeadSelfAttention module +- Variable Bounding as configurable model layers [#13](https://github.com/ecmwf/anemoi-models/issues/13) + +### Changed +- Bugfixes for CI + +### Removed + +## [0.3.0](https://github.com/ecmwf/anemoi-models/compare/0.2.1...0.3.0) - Remapping of (meteorological) Variables +### Added + +- CI workflow to update the changelog on release +- Remapper: Preprocessor for remapping one variable to multiple ones. Includes changes to the data indices since the remapper changes the number of variables. With optional config keywords. ### Changed - 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. -- Bugfixes for CI ### Removed diff --git a/src/anemoi/models/layers/bounding.py b/src/anemoi/models/layers/bounding.py new file mode 100644 index 0000000..3791ff2 --- /dev/null +++ b/src/anemoi/models/layers/bounding.py @@ -0,0 +1,115 @@ +from __future__ import annotations + +from abc import ABC +from abc import abstractmethod + +import torch +from torch import nn + +from anemoi.models.data_indices.tensor import InputTensorIndex + + +class BaseBounding(nn.Module, ABC): + """Abstract base class for bounding strategies. + + This class defines an interface for bounding strategies which are used to apply a specific + restriction to the predictions of a model. + """ + + def __init__( + self, + *, + variables: list[str], + name_to_index: dict, + ) -> None: + super().__init__() + + self.name_to_index = name_to_index + self.variables = variables + self.data_index = self._create_index(variables=self.variables) + + def _create_index(self, variables: list[str]) -> InputTensorIndex: + return InputTensorIndex(includes=variables, excludes=[], name_to_index=self.name_to_index)._only + + @abstractmethod + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Applies the bounding to the predictions. + + Parameters + ---------- + x : torch.Tensor + The tensor containing the predictions that will be bounded. + + Returns + ------- + torch.Tensor + A tensor with the bounding applied. + """ + pass + + +class ReluBounding(BaseBounding): + """Initializes the bounding with a ReLU activation / zero clamping.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x[..., self.data_index] = torch.nn.functional.relu(x[..., self.data_index]) + return x + + +class HardtanhBounding(BaseBounding): + """Initializes the bounding with specified minimum and maximum values for bounding. + + Parameters + ---------- + variables : list[str] + A list of strings representing the variables that will be bounded. + name_to_index : dict + A dictionary mapping the variable names to their corresponding indices. + min_val : float + The minimum value for the HardTanh activation. + max_val : float + The maximum value for the HardTanh activation. + """ + + def __init__(self, *, variables: list[str], name_to_index: dict, min_val: float, max_val: float) -> None: + super().__init__(variables=variables, name_to_index=name_to_index) + self.min_val = min_val + self.max_val = max_val + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x[..., self.data_index] = torch.nn.functional.hardtanh( + x[..., self.data_index], min_val=self.min_val, max_val=self.max_val + ) + return x + + +class FractionBounding(HardtanhBounding): + """Initializes the FractionBounding with specified parameters. + + Parameters + ---------- + variables : list[str] + A list of strings representing the variables that will be bounded. + name_to_index : dict + A dictionary mapping the variable names to their corresponding indices. + min_val : float + The minimum value for the HardTanh activation. + max_val : float + The maximum value for the HardTanh activation. + total_var : str + A string representing a variable from which a secondary variable is derived. For + example, in the case of convective precipitation (Cp), total_var = Tp (total precipitation). + """ + + def __init__( + self, *, variables: list[str], name_to_index: dict, min_val: float, max_val: float, total_var: str + ) -> None: + super().__init__(variables=variables, name_to_index=name_to_index, min_val=min_val, max_val=max_val) + self.total_variable = self._create_index(variables=[total_var]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Apply the HardTanh bounding to the data_index variables + x = super().forward(x) + # Calculate the fraction of the total variable + x[..., self.data_index] *= x[..., self.total_variable] + return x diff --git a/src/anemoi/models/models/encoder_processor_decoder.py b/src/anemoi/models/models/encoder_processor_decoder.py index 3414dc5..b043b0c 100644 --- a/src/anemoi/models/models/encoder_processor_decoder.py +++ b/src/anemoi/models/models/encoder_processor_decoder.py @@ -67,6 +67,8 @@ def __init__( self._register_latlon("data", self._graph_name_data) self._register_latlon("hidden", self._graph_name_hidden) + self.data_indices = data_indices + self.num_channels = config.model.num_channels input_dim = self.multi_step * self.num_input_channels + self.latlons_data.shape[1] + self.trainable_data_size @@ -103,6 +105,14 @@ def __init__( dst_grid_size=self._data_grid_size, ) + # Instantiation of model output bounding functions (e.g., to ensure outputs like TP are positive definite) + self.boundings = nn.ModuleList( + [ + instantiate(cfg, name_to_index=self.data_indices.model.output.name_to_index) + for cfg in getattr(config.model, "bounding", []) + ] + ) + def _calculate_shapes_and_indices(self, data_indices: dict) -> None: self.num_input_channels = len(data_indices.internal_model.input) self.num_output_channels = len(data_indices.internal_model.output) @@ -251,4 +261,9 @@ def forward(self, x: Tensor, model_comm_group: Optional[ProcessGroup] = None) -> # residual connection (just for the prognostic variables) x_out[..., self._internal_output_idx] += x[:, -1, :, :, self._internal_input_idx] + + for bounding in self.boundings: + # bounding performed in the order specified in the config file + x_out = bounding(x_out) + return x_out diff --git a/src/anemoi/models/preprocessing/__init__.py b/src/anemoi/models/preprocessing/__init__.py index 53017fb..cc2cb4f 100644 --- a/src/anemoi/models/preprocessing/__init__.py +++ b/src/anemoi/models/preprocessing/__init__.py @@ -38,7 +38,23 @@ def __init__( Data indices for input and output variables statistics : dict Data statistics dictionary + data_indices : dict + Data indices for input and output variables + + Attributes + ---------- + default : str + Default method for variables not specified in the config + method_config : dict + Dictionary of the methods with lists of variables + methods : dict + Dictionary of the variables with methods + data_indices : IndexCollection + Data indices for input and output variables + remap : dict + Dictionary of the variables with remapped names in the config """ + super().__init__() self.default, self.method_config = self._process_config(config) @@ -47,8 +63,10 @@ def __init__( self.data_indices = data_indices def _process_config(self, config): + _special_keys = ["default", "remap"] # Keys that do not contain a list of variables in a preprocessing method. default = config.get("default", "none") - method_config = {k: v for k, v in config.items() if k != "default" and v is not None and v != "none"} + self.remap = config.get("remap", {}) + method_config = {k: v for k, v in config.items() if k not in _special_keys and v is not None and v != "none"} if not method_config: LOGGER.warning( diff --git a/src/anemoi/models/preprocessing/normalizer.py b/src/anemoi/models/preprocessing/normalizer.py index bc75466..ee6a4f5 100644 --- a/src/anemoi/models/preprocessing/normalizer.py +++ b/src/anemoi/models/preprocessing/normalizer.py @@ -49,6 +49,16 @@ def __init__( mean = statistics["mean"] stdev = statistics["stdev"] + # Optionally reuse statistic of one variable for another variable + statistics_remap = {} + for remap, source in self.remap.items(): + idx_src, idx_remap = name_to_index_training_input[source], name_to_index_training_input[remap] + statistics_remap[idx_remap] = (minimum[idx_src], maximum[idx_src], mean[idx_src], stdev[idx_src]) + + # Two-step to avoid overwriting the original statistics in the loop (this reduces dependence on order) + for idx, new_stats in statistics_remap.items(): + minimum[idx], maximum[idx], mean[idx], stdev[idx] = new_stats + self._validate_normalization_inputs(name_to_index_training_input, minimum, maximum, mean, stdev) _norm_add = np.zeros((minimum.size,), dtype=np.float32) @@ -56,6 +66,7 @@ def __init__( for name, i in name_to_index_training_input.items(): method = self.methods.get(name, self.default) + if method == "mean-std": LOGGER.debug(f"Normalizing: {name} is mean-std-normalised.") if stdev[i] < (mean[i] * 1e-6): @@ -63,6 +74,13 @@ def __init__( _norm_mul[i] = 1 / stdev[i] _norm_add[i] = -mean[i] / stdev[i] + elif method == "std": + LOGGER.debug(f"Normalizing: {name} is std-normalised.") + if stdev[i] < (mean[i] * 1e-6): + warnings.warn(f"Normalizing: the field seems to have only one value {mean[i]}") + _norm_mul[i] = 1 / stdev[i] + _norm_add[i] = 0 + elif method == "min-max": LOGGER.debug(f"Normalizing: {name} is min-max-normalised to [0, 1].") x = maximum[i] - minimum[i] @@ -92,16 +110,20 @@ def _validate_normalization_inputs(self, name_to_index_training_input: dict, min f"Error parsing methods in InputNormalizer methods ({len(self.methods)}) " f"and entries in config ({sum(len(v) for v in self.method_config)}) do not match." ) + + # Check that all sizes align n = minimum.size assert maximum.size == n, (maximum.size, n) assert mean.size == n, (mean.size, n) assert stdev.size == n, (stdev.size, n) + # Check for typos in method config assert isinstance(self.methods, dict) for name, method in self.methods.items(): assert name in name_to_index_training_input, f"{name} is not a valid variable name" assert method in [ "mean-std", + "std", # "robust", "min-max", "max", diff --git a/tests/layers/test_bounding.py b/tests/layers/test_bounding.py new file mode 100644 index 0000000..87619cd --- /dev/null +++ b/tests/layers/test_bounding.py @@ -0,0 +1,92 @@ +import pytest +import torch +from anemoi.utils.config import DotDict +from hydra.utils import instantiate + +from anemoi.models.layers.bounding import FractionBounding +from anemoi.models.layers.bounding import HardtanhBounding +from anemoi.models.layers.bounding import ReluBounding + + +@pytest.fixture +def config(): + return DotDict({"variables": ["var1", "var2"], "total_var": "total_var"}) + + +@pytest.fixture +def name_to_index(): + return {"var1": 0, "var2": 1, "total_var": 2} + + +@pytest.fixture +def input_tensor(): + return torch.tensor([[-1.0, 2.0, 3.0], [4.0, -5.0, 6.0], [0.5, 0.5, 0.5]]) + + +def test_relu_bounding(config, name_to_index, input_tensor): + bounding = ReluBounding(variables=config.variables, name_to_index=name_to_index) + output = bounding(input_tensor.clone()) + expected_output = torch.tensor([[0.0, 2.0, 3.0], [4.0, 0.0, 6.0], [0.5, 0.5, 0.5]]) + assert torch.equal(output, expected_output) + + +def test_hardtanh_bounding(config, name_to_index, input_tensor): + minimum, maximum = -1.0, 1.0 + bounding = HardtanhBounding( + variables=config.variables, name_to_index=name_to_index, min_val=minimum, max_val=maximum + ) + output = bounding(input_tensor.clone()) + expected_output = torch.tensor([[minimum, maximum, 3.0], [maximum, minimum, 6.0], [0.5, 0.5, 0.5]]) + assert torch.equal(output, expected_output) + + +def test_fraction_bounding(config, name_to_index, input_tensor): + bounding = FractionBounding( + variables=config.variables, name_to_index=name_to_index, min_val=0.0, max_val=1.0, total_var=config.total_var + ) + output = bounding(input_tensor.clone()) + expected_output = torch.tensor([[0.0, 3.0, 3.0], [6.0, 0.0, 6.0], [0.25, 0.25, 0.5]]) + + assert torch.equal(output, expected_output) + + +def test_multi_chained_bounding(config, name_to_index, input_tensor): + # Apply Relu first on the first variable only + bounding1 = ReluBounding(variables=config.variables[:-1], name_to_index=name_to_index) + expected_output = torch.tensor([[0.0, 2.0, 3.0], [4.0, -5.0, 6.0], [0.5, 0.5, 0.5]]) + # Check intemediate result + assert torch.equal(bounding1(input_tensor.clone()), expected_output) + minimum, maximum = 0.5, 1.75 + bounding2 = HardtanhBounding( + variables=config.variables, name_to_index=name_to_index, min_val=minimum, max_val=maximum + ) + # Use full chaining on the input tensor + output = bounding2(bounding1(input_tensor.clone())) + # Data with Relu applied first and then Hardtanh + expected_output = torch.tensor([[minimum, maximum, 3.0], [maximum, minimum, 6.0], [0.5, 0.5, 0.5]]) + assert torch.equal(output, expected_output) + + +def test_hydra_instantiate_bounding(config, name_to_index, input_tensor): + layer_definitions = [ + { + "_target_": "anemoi.models.layers.bounding.ReluBounding", + "variables": config.variables, + }, + { + "_target_": "anemoi.models.layers.bounding.HardtanhBounding", + "variables": config.variables, + "min_val": 0.0, + "max_val": 1.0, + }, + { + "_target_": "anemoi.models.layers.bounding.FractionBounding", + "variables": config.variables, + "min_val": 0.0, + "max_val": 1.0, + "total_var": config.total_var, + }, + ] + for layer_definition in layer_definitions: + bounding = instantiate(layer_definition, name_to_index=name_to_index) + bounding(input_tensor.clone()) diff --git a/tests/preprocessing/test_preprocessor_normalizer.py b/tests/preprocessing/test_preprocessor_normalizer.py index cc527e7..8056865 100644 --- a/tests/preprocessing/test_preprocessor_normalizer.py +++ b/tests/preprocessing/test_preprocessor_normalizer.py @@ -40,6 +40,37 @@ def input_normalizer(): return InputNormalizer(config=config.data.normalizer, data_indices=data_indices, statistics=statistics) +@pytest.fixture() +def remap_normalizer(): + config = DictConfig( + { + "diagnostics": {"log": {"code": {"level": "DEBUG"}}}, + "data": { + "normalizer": { + "default": "mean-std", + "remap": {"x": "z", "y": "x"}, + "min-max": ["x"], + "max": ["y"], + "none": ["z"], + "mean-std": ["q"], + }, + "forcing": ["z", "q"], + "diagnostic": ["other"], + "remapped": {}, + }, + }, + ) + statistics = { + "mean": np.array([1.0, 2.0, 3.0, 4.5, 3.0]), + "stdev": np.array([0.5, 0.5, 0.5, 1, 14]), + "minimum": np.array([1.0, 1.0, 1.0, 1.0, 1.0]), + "maximum": np.array([11.0, 10.0, 10.0, 10.0, 10.0]), + } + name_to_index = {"x": 0, "y": 1, "z": 2, "q": 3, "other": 4} + data_indices = IndexCollection(config=config, name_to_index=name_to_index) + return InputNormalizer(config=config.data.normalizer, statistics=statistics, data_indices=data_indices) + + def test_normalizer_not_inplace(input_normalizer) -> None: x = torch.Tensor([[1.0, 2.0, 3.0, 4.0, 5.0], [6.0, 7.0, 8.0, 9.0, 10.0]]) input_normalizer(x, in_place=False) @@ -87,3 +118,15 @@ def test_normalize_inverse_transform(input_normalizer) -> None: assert torch.allclose( input_normalizer.inverse_transform(input_normalizer.transform(x, in_place=False), in_place=False), x ) + + +def test_normalizer_not_inplace_remap(remap_normalizer) -> None: + x = torch.Tensor([[1.0, 2.0, 3.0, 4.0, 5.0], [6.0, 7.0, 8.0, 9.0, 10.0]]) + remap_normalizer(x, in_place=False) + assert torch.allclose(x, torch.Tensor([[1.0, 2.0, 3.0, 4.0, 5.0], [6.0, 7.0, 8.0, 9.0, 10.0]])) + + +def test_normalize_remap(remap_normalizer) -> None: + x = torch.Tensor([[1.0, 2.0, 3.0, 4.0, 5.0], [6.0, 7.0, 8.0, 9.0, 10.0]]) + expected_output = torch.Tensor([[0.0, 2 / 11, 3.0, -0.5, 1 / 7], [5 / 9, 7 / 11, 8.0, 4.5, 0.5]]) + assert torch.allclose(remap_normalizer.transform(x), expected_output) From e608a736fe0fa6d01a8a11a2c96072e116c2877a Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 23 Sep 2024 16:52:14 +0200 Subject: [PATCH 5/6] [Feature] 28 Make models switchable through the config (#45) * feat: make model instantiateable * docs: instantiation explained in changelog * refactor: rename model config object * fix: rename config to model_config * fix: mark non-recursive * docs: changelog --- CHANGELOG.md | 1 + src/anemoi/models/interface/__init__.py | 11 +++++---- .../models/encoder_processor_decoder.py | 24 +++++++++---------- 3 files changed, 20 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7811ef3..5678486 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,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 [#45](https://github.com/ecmwf/anemoi-models/pulls/45) ### Removed diff --git a/src/anemoi/models/interface/__init__.py b/src/anemoi/models/interface/__init__.py index 626940f..aba62a2 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,13 @@ 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, + 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 diff --git a/src/anemoi/models/models/encoder_processor_decoder.py b/src/anemoi/models/models/encoder_processor_decoder.py index b043b0c..aa7e8bb 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() @@ -69,13 +69,13 @@ def __init__( self.data_indices = data_indices - 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, @@ -86,7 +86,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, @@ -95,7 +95,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, @@ -109,7 +109,7 @@ def __init__( self.boundings = nn.ModuleList( [ instantiate(cfg, name_to_index=self.data_indices.model.output.name_to_index) - for cfg in getattr(config.model, "bounding", []) + for cfg in getattr(model_config.model, "bounding", []) ] ) From 6e623b983179d6228d48abce2f1ddf18defc988d Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Tue, 24 Sep 2024 19:04:28 +0200 Subject: [PATCH 6/6] ci: fix publish and PR workflow (#48) --- .github/workflows/python-publish.yml | 40 ++++------------------- .github/workflows/python-pull-request.yml | 19 +++++++++++ 2 files changed, 26 insertions(+), 33 deletions(-) create mode 100644 .github/workflows/python-pull-request.yml diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml index de01bf6..2cb554a 100644 --- a/.github/workflows/python-publish.yml +++ b/.github/workflows/python-publish.yml @@ -4,48 +4,22 @@ name: Upload Python Package on: - - push: {} - release: types: [created] jobs: quality: - name: Code QA - runs-on: ubuntu-latest - steps: - - run: sudo apt-get install -y pandoc # Needed by sphinx for notebooks - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 - with: - python-version: 3.x - - uses: pre-commit/action@v3.0.1 + uses: ecmwf-actions/reusable-workflows/.github/workflows/qa-precommit-run.yml@v2 + with: + skip-hooks: "no-commit-to-branch" checks: strategy: - fail-fast: false matrix: - platform: ["ubuntu-latest", "macos-latest"] - python-version: ["3.10"] - - name: Python ${{ matrix.python-version }} on ${{ matrix.platform }} - runs-on: ${{ matrix.platform }} - - steps: - - uses: actions/checkout@v4 - - - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - - name: Install - run: | - pip install -e .[all,tests] - pip freeze - - - name: Tests - run: pytest + python-version: ["3.9", "3.10"] + uses: ecmwf-actions/reusable-workflows/.github/workflows/qa-pytest-pyproject.yml@v2 + with: + python-version: ${{ matrix.python-version }} deploy: needs: [checks, quality] diff --git a/.github/workflows/python-pull-request.yml b/.github/workflows/python-pull-request.yml new file mode 100644 index 0000000..3488f55 --- /dev/null +++ b/.github/workflows/python-pull-request.yml @@ -0,0 +1,19 @@ +--- +# This workflow will upload a Python Package using Twine when a release is created +# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries + +name: Code Quality checks for PRs + +on: + push: + pull_request: + types: [opened, synchronize, reopened] + +jobs: + quality: + uses: ecmwf-actions/reusable-workflows/.github/workflows/qa-precommit-run.yml@v2 + with: + skip-hooks: "no-commit-to-branch" + + checks: + uses: ecmwf-actions/reusable-workflows/.github/workflows/qa-pytest-pyproject.yml@v2