WIP loss function / forecaster multiple datasets

Author: havardhhaugen

models/src/anemoi/models/models/encoder_processor_decoder.py

@@ -79,6 +79,22 @@ def __init__(
         skip_variables = [var for var in self.input_variables if var in self.output_variables]
         self._internal_input_idx = [self.input_variables.index(var) for var in skip_variables]
         self._internal_output_idx = [self.output_variables.index(var) for var in skip_variables]
+
+        #Fake data_indices to be used in Forecaster
+        self.data_indices = DotDict(
+            {
+                "internal_model": {
+                    "input": {
+                        "full": self.input_variables,
+                        "name_to_index": self.name_to_index_input,
+                    },
+                    "output": {
+                        "full": self.output_variables,
+                        "name_to_index": self.name_to_index_output,
+                    }
+                }
+            }
+        )
         #-------------------------------------------------------------------------#
 
         self.data_indices = data_indices

training/src/anemoi/training/config/model/refactored_graphtransformer.yaml

@@ -1,7 +1,7 @@
 activation: GELU
 num_channels: 1024
 cpu_offload: False
-output_mask: null
+#output_mask: null # moved this to config.training to have one mask per output
 
 model:
   _target_: anemoi.models.models.encoder_processor_decoder.AnemoiModelEncProcDec

training/src/anemoi/training/config/training/refactored_default.yaml

@@ -0,0 +1,129 @@
+# resume or fork a training from a checkpoint last.ckpt or specified in hardware.files.warm_start
+run_id: null
+fork_run_id: null
+transfer_learning: False # activate to perform transfer learning
+load_weights_only: False # only load model weights, do not restore optimiser states etc.
+
+# run in deterministic mode ; slows down
+deterministic: False
+
+# miscellaneous
+precision: 16-mixed
+
+# multistep input
+# 1 = single step scheme, X(t-1) used to predict X(t)
+# k > 1: multistep scheme, uses [X(t-k), X(t-k+1), ... X(t-1)] to predict X(t)
+# Deepmind use k = 2 in their model
+multistep_input: 2
+
+# gradient accumulation across K batches, K >= 1 (if K == 1 then no accumulation)
+# the effective batch size becomes num-devices * batch_size * k
+accum_grad_batches: 1
+
+num_sanity_val_steps: 6
+
+# clipp gradients, 0 : don't clip, default algorithm: norm, alternative: value
+gradient_clip:
+  val: 32.
+  algorithm: value
+
+# stochastic weight averaging
+# https://pytorch.org/blog/stochastic-weight-averaging-in-pytorch/
+swa:
+  enabled: False
+  lr: 1.e-4
+
+# Optimizer settings
+optimizer:
+  zero: False # use ZeroRedundancyOptimizer ; saves memory for larger models
+  kwargs:
+    betas: [0.9, 0.95]
+
+# select model
+model_task: anemoi.training.train.forecaster.GraphForecasterMultiDataset
+
+# select strategy
+strategy:
+  _target_: anemoi.training.distributed.strategy.DDPGroupStrategy
+  num_gpus_per_model: ${hardware.num_gpus_per_model}
+  read_group_size: ${dataloader.read_group_size}
+
+# loss functions
+
+# dynamic rescaling of the loss gradient
+# see https://arxiv.org/pdf/2306.06079.pdf, section 4.3.2
+# don't enable this by default until it's been tested and proven beneficial
+loss_gradient_scaling: False
+
+# length of the "rollout" window (see Keisler's paper)
+rollout:
+  start: 1
+  # increase rollout every n epochs
+  epoch_increment: 0
+  # maximum rollout to use
+  max: 1
+
+# Set max_epochs or max_steps. Training stops at the first limit reached.
+max_epochs: null
+max_steps: 150000
+
+lr:
+  warmup: 1000 # number of warmup iterations
+  rate: 0.625e-4 #local_lr
+  iterations: ${training.max_steps} # NOTE: When max_epochs < max_steps, scheduler will run for max_steps
+  min: 3e-7 #Not scaled by #GPU
+
+output:
+  era5: # Make it easier to set the same key name in different places of the config
+    training_loss:
+      _target_: anemoi.training.losses.mse.WeightedMSELoss
+      scalars: ['variable', 'loss_weights_mask']
+      ignore_nans: True
+    output_mask: null
+    variable_loss_scaling:
+      default: 1
+      pl: 
+        q: 0.6
+        t: 6
+        u: 0.8
+        v: 0.5
+        w: 0.001
+        z: 12
+      sfc:
+        sp: 10
+        10u: 0.1
+        10v: 0.1
+        2d: 0.5
+        tp: 0.025
+        cp: 0.0025
+
+    pressure_level_scaler:
+      _target_: anemoi.training.data.scaling.ReluPressureLevelScaler
+      minimum: 0.2
+      slope: 0.001
+
+    node_loss_weights:
+      _target_: anemoi.training.losses.nodeweights.GraphNodeAttribute
+      target_nodes: ${graph.data}
+      node_attribute: area_weight
+
+    validation_metrics:
+      - _target_: anemoi.training.losses.mse.WeightedMSELoss
+        scalars: []
+        ignore_nans: True
+    scale_validation_metrics:
+      scalars_to_apply: ['variable']
+      metrics:
+        - 'all'
+    
+    metrics:
+      - z_500
+      - t_850
+      - u_850
+      - v_850
+
+
+submodules_to_freeze: []
+    
+
+

training/src/anemoi/training/losses/dict.py

@@ -0,0 +1,34 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+import torch.nn as nn
+from torch import Tensor
+
+class DictLoss(nn.Module):
+    """Wrapper for a dictionary of loss-fuctions that operate on different outputs."""
+
+    def __init__(
+        self, 
+        loss_dict: nn.ModuleDict,
+    ) -> None:
+        super().__init__()
+        self.loss_dict = loss_dict
+        self.outputs = list(loss_dict.keys())
+
+    def forward(
+        self,
+        pred: dict[str, Tensor],
+        target: dict[str, Tensor],
+        squash: bool = True, # TODO Generalise this per output?
+        ) -> dict[str, Tensor]:
+        out = {}
+        for output, loss in self.loss_dict.items():
+            out[output] = loss(pred[output], target[output], squash)
+
+        return out

training/src/anemoi/training/train/forecaster/__init__.py

@@ -10,5 +10,6 @@
 from .ensforecaster import GraphEnsForecaster
 from .forecaster import GraphForecaster
 from .interpolator import GraphInterpolator
+from .forecaster_multiple_datasets import GraphForecasterMultiDataset
 
-__all__ = ["GraphEnsForecaster", "GraphForecaster", "GraphInterpolator"]
+__all__ = ["GraphEnsForecaster", "GraphForecaster", "GraphInterpolator", "GraphForecasterMultiDataset"]