diff --git a/aurora/__init__.py b/aurora/__init__.py index 9403fb7f..8f414aac 100644 --- a/aurora/__init__.py +++ b/aurora/__init__.py @@ -14,7 +14,7 @@ AuroraV1p5Ensemble, AuroraWave, ) -from aurora.rollout import rollout +from aurora.rollout import rollout, rollout_ensemble from aurora.tracker import Tracker __all__ = [ @@ -32,5 +32,6 @@ "Metadata", "insolation", "rollout", + "rollout_ensemble", "Tracker", ] diff --git a/aurora/batch.py b/aurora/batch.py index d0616d7b..dca332df 100644 --- a/aurora/batch.py +++ b/aurora/batch.py @@ -311,6 +311,37 @@ def from_netcdf(cls, path: str | Path) -> "Batch": ) +def _tile_batch(batch: Batch, n: int) -> Batch: + """Tile `batch` along the batch dimension `n` times. + + Not part of the public `Batch` API. Used only by `aurora.Aurora.forward` and + `aurora.rollout.rollout` to run `n` ensemble members as a single fused computation. The + tiled batch dimension is an internal implementation detail and must be undone with + `_split_batch` before any result derived from it is returned to a caller. + """ + return dataclasses.replace( + batch, + surf_vars={k: v.repeat(n, *([1] * (v.dim() - 1))) for k, v in batch.surf_vars.items()}, + atmos_vars={k: v.repeat(n, *([1] * (v.dim() - 1))) for k, v in batch.atmos_vars.items()}, + metadata=dataclasses.replace(batch.metadata, time=batch.metadata.time * n), + ) + + +def _split_batch(batch: Batch, n: int) -> list[Batch]: + """Undo `_tile_batch`, splitting a tiled batch back into `n` standard-shaped batches.""" + b = next(iter(batch.surf_vars.values())).shape[0] // n + time = batch.metadata.time + return [ + dataclasses.replace( + batch, + surf_vars={k: v[m * b : (m + 1) * b] for k, v in batch.surf_vars.items()}, + atmos_vars={k: v[m * b : (m + 1) * b] for k, v in batch.atmos_vars.items()}, + metadata=dataclasses.replace(batch.metadata, time=time[m * b : (m + 1) * b]), + ) + for m in range(n) + ] + + def _np(x: torch.Tensor) -> np.ndarray: return x.detach().cpu().numpy() diff --git a/aurora/model/aurora.py b/aurora/model/aurora.py index 059dad24..7886955a 100644 --- a/aurora/model/aurora.py +++ b/aurora/model/aurora.py @@ -14,7 +14,7 @@ apply_activation_checkpointing, ) -from aurora.batch import Batch +from aurora.batch import Batch, _split_batch, _tile_batch from aurora.insolation import insolation from aurora.model.compat import ( _adapt_checkpoint_air_pollution, @@ -103,6 +103,7 @@ def __init__( clamp_at_first_step: bool = False, simulate_indexing_bug: bool = False, stochastic: bool = False, + num_ensemble_members: int = 1, use_updated_lead_time_embedding: bool = False, variable_lead_time: bool = False, rollout_input_clipping: Optional[dict[str, dict[str, Optional[float]]]] = None, @@ -200,6 +201,17 @@ def __init__( to the original implementation. Defaults to `False`. stochastic (bool, optional): If `True`, enable stochastic mode with noise injection. Defaults to `False`. + num_ensemble_members (int, optional): Number of ensemble members to produce + *internally* on every call to :meth:`forward_ensemble`, as an alternative to + looping over separate :meth:`forward` calls and combining the results externally + yourself (which remains perfectly valid, e.g. if you need more control over how + members are seeded or combined). When set to a value greater than `1`, the batch + is tiled `num_ensemble_members` times internally and run through the model in a + single, fully-batched pass, which is far more efficient on a GPU than looping. + This is most useful in combination with `stochastic=True`, since every tiled copy + then receives independent noise. When greater than `1`, plain :meth:`forward` + raises, since it can only ever return a single `Batch`; use + :meth:`forward_ensemble` instead. Defaults to `1`, i.e. no internal ensembling. use_updated_lead_time_embedding (bool, optional): Whether to use the updated lead time embedding with a minimum wavelength of 2 hours. Defaults to `False`. variable_lead_time (bool, optional): If `True`, use per-sample lead times passed @@ -236,6 +248,10 @@ def __init__( self.output_only_surf_vars = output_only_surf_vars self.output_only_atmos_vars = output_only_atmos_vars + if num_ensemble_members < 1: + raise ValueError("`num_ensemble_members` must be at least `1`.") + self.num_ensemble_members = num_ensemble_members + if self.surf_stats: warnings.warn( f"The normalisation statics for the following surface-level variables are manually " @@ -284,6 +300,14 @@ def __init__( use_updated_lead_time_embedding=use_updated_lead_time_embedding, ) + if num_ensemble_members > 1 and not self.backbone.stochastic: + warnings.warn( + f"`num_ensemble_members={num_ensemble_members}` was requested, but `stochastic=" + f"False`, so the model has no source of randomness. All ensemble members will be " + f"identical.", + stacklevel=2, + ) + self.decoder = Perceiver3DDecoder( surf_vars=surf_vars, atmos_vars=atmos_vars, @@ -349,7 +373,47 @@ def forward(self, batch: Batch, lead_times: Optional[torch.Tensor] = None) -> Ba `variable_lead_time=True`. Ignored otherwise. Returns: - :class:`Batch`: Prediction for the batch. + :class:`Batch`: Prediction for `batch`. + """ + if self.num_ensemble_members > 1: + raise RuntimeError( + f"This model was constructed with `num_ensemble_members=" + f"{self.num_ensemble_members}`. Use `forward_ensemble` instead of `forward` to " + f"obtain all ensemble members." + ) + return self._forward_impl(batch, lead_times) + + def forward_ensemble( + self, batch: Batch, lead_times: Optional[torch.Tensor] = None + ) -> list[Batch]: + """Forward pass producing all `self.num_ensemble_members` ensemble members internally. + + All members are computed internally as a single fused pass through the + encoder/backbone/decoder, rather than looping over separate `forward` calls and combining + the results externally yourself (which remains equally valid and unaffected). This is most + useful in combination with `stochastic=True`, since every internally-tiled copy then + receives independent noise. + + Args: + batch (:class:`aurora.Batch`): Batch to run the model on. + lead_times (:class:`torch.Tensor`, optional): Per-sample lead times of shape + `(batch,)` in hours. Required when the model was configured with + `variable_lead_time=True`. Ignored otherwise. + + Returns: + list[:class:`Batch`]: A list of `self.num_ensemble_members` standard-shaped `Batch`\\ + s, one per ensemble member, each with the same batch dimension as `batch`. + """ + pred = self._forward_impl(batch, lead_times) + return _split_batch(pred, self.num_ensemble_members) + + def _forward_impl(self, batch: Batch, lead_times: Optional[torch.Tensor] = None) -> Batch: + """Shared implementation for `forward` and `forward_ensemble`. + + Internally tiles `batch` by `self.num_ensemble_members` before running it through the + encoder/backbone/decoder as a single fused batch, when greater than `1`. The tiled batch + dimension is a private implementation detail: `forward` forbids it (see above) and + `forward_ensemble` splits it back apart before returning. """ batch = self.batch_transform_hook(batch) @@ -361,6 +425,11 @@ def forward(self, batch: Batch, lead_times: Optional[torch.Tensor] = None) -> Ba batch = batch.crop(patch_size=self.patch_size) batch = batch.to(p.device) + if self.num_ensemble_members > 1: + if lead_times is not None: + lead_times = lead_times.repeat(self.num_ensemble_members) + batch = _tile_batch(batch, self.num_ensemble_members) + H, W = batch.spatial_shape patch_res = ( self.encoder.latent_levels, diff --git a/aurora/rollout.py b/aurora/rollout.py index 3a84ce6b..af166a03 100644 --- a/aurora/rollout.py +++ b/aurora/rollout.py @@ -6,10 +6,10 @@ import torch -from aurora.batch import Batch +from aurora.batch import Batch, _split_batch, _tile_batch from aurora.model.aurora import Aurora -__all__ = ["rollout"] +__all__ = ["rollout", "rollout_ensemble"] def _make_lead_time_tensor(batch: Batch, lead_time_hours: float) -> torch.Tensor: @@ -148,3 +148,44 @@ def rollout( # Disable noise accumulation after roll-out is complete, in case the model will be used for # normal inference or training afterwards. model.set_noise_accumulation(n=0) + + +def rollout_ensemble( + model: Aurora, + batch: Batch, + steps: int, + fine_lead_times: Optional[Sequence[float]] = None, + use_noise_accumulation: bool = True, + apply_rollout_input_clipping: bool = True, +) -> Generator[list[Batch], None, None]: + """Like `rollout`, but produces `model.num_ensemble_members` ensemble members internally on + every step, as a single fused pass, instead of `rollout` yielding one `Batch` per step. + + All arguments are identical to `rollout`; see there for details. + + Yields: + list[:class:`aurora.Batch`]: A list of `model.num_ensemble_members` standard-shaped + `Batch`\\ s after every (sub-)step, one per ensemble member; see + :meth:`aurora.Aurora.forward_ensemble`. + """ + num_ensemble_members = model.num_ensemble_members + + # Tile the batch once up front, then temporarily disable further expansion so it isn't + # repeated by `_forward_impl` on every step. The tiled representation is kept purely internal + # to this loop: every `pred` yielded to the caller is split back into standard-shaped batches. + batch = _tile_batch(batch, num_ensemble_members) + model.num_ensemble_members = 1 + try: + for pred in rollout( + model, + batch, + steps, + fine_lead_times=fine_lead_times, + use_noise_accumulation=use_noise_accumulation, + apply_rollout_input_clipping=apply_rollout_input_clipping, + ): + yield _split_batch(pred, num_ensemble_members) + finally: + # Restore the model's ensemble configuration, whether the roll-out ran to completion or + # was abandoned early. + model.num_ensemble_members = num_ensemble_members diff --git a/tests/v1p5/test_ensemble.py b/tests/v1p5/test_ensemble.py new file mode 100644 index 00000000..db7a3755 --- /dev/null +++ b/tests/v1p5/test_ensemble.py @@ -0,0 +1,252 @@ +"""Copyright (c) Microsoft Corporation. Licensed under the MIT license. + +Tests for internal ensemble members (`num_ensemble_members`). +""" + +import warnings +from datetime import datetime + +import pytest +import torch + +from ._helpers import _OUTPUT_ONLY_SURF, _SURF_VARS, _make_batch, _make_small_v1p5 +from aurora import Aurora, Batch, Metadata, rollout, rollout_ensemble +from aurora.batch import _split_batch, _tile_batch +from aurora.model.film import AdaptiveLayerNorm + + +def _unzero_adaptive_layer_norms(model: Aurora, std: float = 0.1) -> None: + """Nudge every `AdaptiveLayerNorm`'s modulation away from its zero initialisation. + + At construction, `AdaptiveLayerNorm.ln_modulation` is exactly zero-initialised (the + `adaLN-Zero` trick), which makes a freshly-built, untrained model exactly insensitive to its + conditioning signal `c` -- which is what carries the ensemble noise. Without this, no output + difference a test observes between ensemble members can be attributed to noise, since noise + provably has zero effect on such a model. + """ + for m in model.modules(): + if isinstance(m, AdaptiveLayerNorm): + with torch.no_grad(): + m.ln_modulation[-1].weight.normal_(std=std) + m.ln_modulation[-1].bias.normal_(std=std) + + +def _make_ensemble_test_batch(b: int = 2) -> Batch: + """A small batch with a configurable batch size `b`, used to test tiling/splitting.""" + h, w = 8, 8 + return Batch( + surf_vars={"2t": torch.randn(b, 2, h, w)}, + static_vars={"lsm": torch.randn(h, w)}, + atmos_vars={"z": torch.randn(b, 2, 2, h, w)}, + metadata=Metadata( + lat=torch.linspace(90, -90, h), + lon=torch.linspace(0, 360, w + 1)[:-1], + time=tuple(datetime(2023, 6, 15, i, 0) for i in range(b)), + atmos_levels=(500, 850), + ), + ) + + +def test_tile_and_split_batch_roundtrip(): + b, n = 2, 3 + batch = _make_ensemble_test_batch(b) + + tiled = _tile_batch(batch, n) + + v = tiled.surf_vars["2t"] + assert v.shape[0] == n * b + for m in range(n): + torch.testing.assert_close(v[m * b : (m + 1) * b], batch.surf_vars["2t"]) + + v = tiled.atmos_vars["z"] + assert v.shape[0] == n * b + for m in range(n): + torch.testing.assert_close(v[m * b : (m + 1) * b], batch.atmos_vars["z"]) + + assert len(tiled.metadata.time) == n * b + for m in range(n): + assert tiled.metadata.time[m * b : (m + 1) * b] == batch.metadata.time + + # Static variables have no batch dimension and are untouched. + torch.testing.assert_close(tiled.static_vars["lsm"], batch.static_vars["lsm"]) + + # Splitting undoes the tiling: every member is identical to the original, standard-shaped + # batch (tiling itself introduces no randomness). + members = _split_batch(tiled, n) + assert len(members) == n + for member in members: + torch.testing.assert_close(member.surf_vars["2t"], batch.surf_vars["2t"]) + torch.testing.assert_close(member.atmos_vars["z"], batch.atmos_vars["z"]) + assert member.metadata.time == batch.metadata.time + + +def test_num_ensemble_members_must_be_positive(): + with pytest.raises(ValueError, match="num_ensemble_members"): + _make_small_v1p5(num_ensemble_members=0) + + +def test_num_ensemble_members_warns_without_stochastic(): + with pytest.warns(UserWarning, match="stochastic"): + _make_small_v1p5(num_ensemble_members=2, stochastic=False) + + +def test_num_ensemble_members_no_warning_with_stochastic(): + with warnings.catch_warnings(): + warnings.simplefilter("error") + _make_small_v1p5(num_ensemble_members=2, stochastic=True) + + +def test_forward_returns_batch_when_num_ensemble_members_one(): + model = _make_small_v1p5() + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + + with torch.inference_mode(): + pred = model.forward(batch, lead_times=torch.full((1,), 6.0)) + + assert isinstance(pred, Batch) + + +def test_forward_raises_when_num_ensemble_members_greater_than_one(): + model = _make_small_v1p5(stochastic=True, num_ensemble_members=3) + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + + with pytest.raises(RuntimeError, match="forward_ensemble"): + model.forward(batch, lead_times=torch.full((1,), 6.0)) + + +def test_forward_ensemble_returns_list_of_standard_shaped_batches(): + n = 3 + model = _make_small_v1p5(stochastic=True, num_ensemble_members=n) + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + b = next(iter(batch.surf_vars.values())).shape[0] + + with torch.inference_mode(): + pred = model.forward_ensemble(batch, lead_times=torch.full((b,), 6.0)) + + assert isinstance(pred, list) + assert len(pred) == n + for member in pred: + assert isinstance(member, Batch) + for v in member.surf_vars.values(): + assert v.shape[0] == b + for v in member.static_vars.values(): + # Static variables have no batch dimension. + assert v.dim() == 2 + + +def test_forward_ensemble_members_differ_when_stochastic(): + n = 3 + torch.manual_seed(0) + model = _make_small_v1p5(stochastic=True, num_ensemble_members=n) + # Un-zero the modulation so noise has a real, appreciable effect (see helper docstring); + # otherwise this test cannot distinguish genuine noise sensitivity from incidental + # floating-point batching noise + # (see `test_forward_ensemble_members_identical_without_stochastic`). + _unzero_adaptive_layer_norms(model) + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + b = next(iter(batch.surf_vars.values())).shape[0] + + with torch.inference_mode(): + pred = model.forward_ensemble(batch, lead_times=torch.full((b,), 6.0)) + + # Threshold well above the ~1e-3 floating-point batching floor established in + # `test_forward_ensemble_members_identical_without_stochastic`, so a pass here can only be + # explained by the injected noise actually differing per member, not incidental rounding. + for i in range(n): + for j in range(i + 1, n): + diff = (pred[i].surf_vars["2t"] - pred[j].surf_vars["2t"]).abs().max() + assert diff > 1e-2 + + +def test_forward_ensemble_members_identical_without_stochastic(): + n = 3 + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + model = _make_small_v1p5(stochastic=False, num_ensemble_members=n) + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + b = next(iter(batch.surf_vars.values())).shape[0] + + with torch.inference_mode(): + pred = model.forward_ensemble(batch, lead_times=torch.full((b,), 6.0)) + + for m in range(1, n): + # Loose tolerance: floating-point ops (e.g. batched matmul/softmax reductions) are not + # strictly invariant to how many other (tiled) rows share the batch, so bitwise equality + # isn't guaranteed even though the members are mathematically identical computations. + torch.testing.assert_close( + pred[0].surf_vars["2t"], pred[m].surf_vars["2t"], atol=1e-3, rtol=1e-3 + ) + + +def test_rollout_ensemble_yields_list_of_standard_shaped_batches_across_steps(): + n = 2 + model = _make_small_v1p5(stochastic=True, num_ensemble_members=n) + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + b = next(iter(batch.surf_vars.values())).shape[0] + + with torch.inference_mode(): + preds = list(rollout_ensemble(model, batch, steps=3)) + + assert len(preds) == 3 + for step_pred in preds: + assert len(step_pred) == n + for member in step_pred: + for v in member.surf_vars.values(): + assert v.shape[0] == b + + # The model's ensemble configuration is restored after the roll-out completes. + assert model.num_ensemble_members == n + + +def test_rollout_raises_when_num_ensemble_members_greater_than_one(): + model = _make_small_v1p5(stochastic=True, num_ensemble_members=2) + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + + with torch.inference_mode(), pytest.raises(RuntimeError, match="forward_ensemble"): + next(rollout(model, batch, steps=1)) + + +def test_rollout_ensemble_restores_num_ensemble_members_on_early_close(): + n = 2 + model = _make_small_v1p5(stochastic=True, num_ensemble_members=n) + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + + with torch.inference_mode(): + gen = rollout_ensemble(model, batch, steps=5) + next(gen) + gen.close() + + assert model.num_ensemble_members == n + + +def test_rollout_ensemble_num_ensemble_members_one_still_works(): + model = _make_small_v1p5() + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + b = next(iter(batch.surf_vars.values())).shape[0] + + with torch.inference_mode(): + preds = list(rollout_ensemble(model, batch, steps=2)) + + for step_pred in preds: + assert len(step_pred) == 1 + for v in step_pred[0].surf_vars.values(): + assert v.shape[0] == b + assert model.num_ensemble_members == 1