diff --git a/aurora/__init__.py b/aurora/__init__.py index 9403fb7f..5ca4f317 100644 --- a/aurora/__init__.py +++ b/aurora/__init__.py @@ -1,6 +1,6 @@ """Copyright (c) Microsoft Corporation. Licensed under the MIT license.""" -from aurora.batch import Batch, Metadata +from aurora.batch import Batch, Metadata, split_batch, tile_batch from aurora.insolation import insolation from aurora.model.aurora import ( Aurora, @@ -14,7 +14,7 @@ AuroraV1p5Ensemble, AuroraWave, ) -from aurora.rollout import rollout +from aurora.rollout import rollout, rollout_ensemble from aurora.tracker import Tracker __all__ = [ @@ -30,7 +30,10 @@ "AuroraV1p5Ensemble", "Batch", "Metadata", + "tile_batch", + "split_batch", "insolation", "rollout", + "rollout_ensemble", "Tracker", ] diff --git a/aurora/batch.py b/aurora/batch.py index d0616d7b..00ed6b80 100644 --- a/aurora/batch.py +++ b/aurora/batch.py @@ -17,7 +17,7 @@ unnormalise_surf_var, ) -__all__ = ["Metadata", "Batch"] +__all__ = ["Metadata", "Batch", "tile_batch", "split_batch"] @dataclasses.dataclass @@ -311,6 +311,50 @@ def from_netcdf(cls, path: str | Path) -> "Batch": ) +def tile_batch(batch: Batch, n: int) -> Batch: + """Tile `batch` along the batch dimension `n` times. + + Used to run `n` ensemble members as a single fused computation. + Results derived from the tiling should be undone with `split_batch`. + + Args: + batch (:class:`aurora.Batch`): The batch to tile. + n (int): Number of times to tile. + + Returns: + :class:`aurora.Batch`: `batch` tiled `n` times along the batch dimension. + """ + 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. + + Args: + batch (:class:`aurora.Batch`): The tiled batch to split. + n (int): Number of batches `batch` was tiled into. + + Returns: + list[:class:`aurora.Batch`]: `batch` split 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/rollout.py b/aurora/rollout.py index 3a84ce6b..002ab4fd 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,43 @@ 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, + num_ensemble_members: int, + fine_lead_times: Optional[Sequence[float]] = None, + use_noise_accumulation: bool = True, + apply_rollout_input_clipping: bool = True, +) -> Generator[list[Batch], None, None]: + """Perform a roll-out for `num_ensemble_members` ensemble members simultaneously. + + The members are computed in one pass by repeating `batch` along the batch dimension, so every + member receives independent noise. All other arguments are as for :func:`rollout`. + + Args: + num_ensemble_members (int): Number of ensemble members. + + Yields: + list[:class:`aurora.Batch`]: After every (sub-)step, one prediction per ensemble member, + each with the batch size of `batch`. + """ + if num_ensemble_members < 1: + raise ValueError( + f"`num_ensemble_members` must be at least `1`, but is `{num_ensemble_members}`." + ) + if not model.backbone.stochastic: + raise ValueError("`rollout_ensemble` requires a stochastic model.") + + batch = tile_batch(batch, num_ensemble_members) + 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) diff --git a/docs/api.rst b/docs/api.rst index 7540ae0d..236fabcf 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -9,10 +9,15 @@ Batch .. autoclass:: aurora.Metadata :members: +.. autofunction:: aurora.tile_batch + +.. autofunction:: aurora.split_batch + Roll-Outs --------- -.. autoclass:: aurora.rollout - :members: +.. autofunction:: aurora.rollout + +.. autofunction:: aurora.rollout_ensemble Tropical Cyclone Tracking ------------------------- diff --git a/docs/models.md b/docs/models.md index 874edddb..3d605ca4 100644 --- a/docs/models.md +++ b/docs/models.md @@ -449,3 +449,25 @@ When using `rollout` with `fine_lead_times`, noise accumulation is enabled by de smoother intra-step transitions while using independent effective noise between main steps, matching the training regimen. Set `use_noise_accumulation=False` to draw independent noise at each sub-step instead, though this is not recommended. + +### Internal Ensembling + +For a stochastic model, such as `AuroraV1p5Ensemble`, you can use `rollout_ensemble` to +generate an autoregressive forecast with multiple ensemble members at once: + +```python +from aurora import rollout_ensemble + +with torch.inference_mode(): + preds = [ + [member.to("cpu") for member in members] + for members in rollout_ensemble(model, batch, steps=10, num_ensemble_members=4) + ] +``` + +Alternatively, you can produce an ensemble forecast by calling `rollout` multiple times, +once for each ensemble member. +`rollout_ensemble` can be faster, at the cost of requiring more memory. + +Internally, `rollout_ensemble` use `aurora.tile_batch` and `aurora.split_batch`. +These functions can also be used for fused forward passes. diff --git a/tests/v1p5/_helpers.py b/tests/v1p5/_helpers.py index 4ddd7a65..575f357b 100644 --- a/tests/v1p5/_helpers.py +++ b/tests/v1p5/_helpers.py @@ -26,16 +26,17 @@ def _make_batch( surf_vars: tuple[str, ...] = _SURF_VARS, static_vars: tuple[str, ...] = _STATIC_VARS, atmos_vars: tuple[str, ...] = _ATMOS_VARS, + batch_size: int = BATCH, ) -> Batch: """Create a minimal synthetic batch.""" return Batch( - surf_vars={k: torch.randn(BATCH, HISTORY, H, W) for k in surf_vars}, + surf_vars={k: torch.randn(batch_size, HISTORY, H, W) for k in surf_vars}, static_vars={k: torch.randn(H, W) for k in static_vars}, - atmos_vars={k: torch.randn(BATCH, HISTORY, N_LEVELS, H, W) for k in atmos_vars}, + atmos_vars={k: torch.randn(batch_size, HISTORY, N_LEVELS, H, W) for k in atmos_vars}, metadata=Metadata( lat=torch.linspace(90, -90, H), lon=torch.linspace(0, 360, W + 1)[:-1], - time=(datetime(2023, 6, 15, 12, 0),), + time=(datetime(2023, 6, 15, 12, 0),) * batch_size, atmos_levels=(100, 250, 500, 850), ), ) diff --git a/tests/v1p5/test_ensemble.py b/tests/v1p5/test_ensemble.py new file mode 100644 index 00000000..9a940736 --- /dev/null +++ b/tests/v1p5/test_ensemble.py @@ -0,0 +1,113 @@ +"""Copyright (c) Microsoft Corporation. Licensed under the MIT license. + +Tests for ensemble rollout and tiling utility functions. +""" + +import itertools + +import pytest +import torch + +from ._helpers import _OUTPUT_ONLY_SURF, _SURF_VARS, BATCH, _make_batch, _make_small_v1p5 +from aurora import Aurora, 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. + + A freshly constructed model is exactly insensitive to its conditioning signal, and hence to + the injected noise, because `AdaptiveLayerNorm.ln_modulation` is zero-initialised. + """ + 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 test_tile_and_split_batch_roundtrip(): + b, n = 2, 3 + batch = _make_batch(batch_size=b) + + tiled = tile_batch(batch, n) + assert len(tiled.metadata.time) == n * b + for v in (*tiled.surf_vars.values(), *tiled.atmos_vars.values()): + assert v.shape[0] == n * b + + members = split_batch(tiled, n) + assert len(members) == n + for member in members: + assert member.metadata.time == batch.metadata.time + for k, v in member.surf_vars.items(): + torch.testing.assert_close(v, batch.surf_vars[k]) + for k, v in member.atmos_vars.items(): + torch.testing.assert_close(v, batch.atmos_vars[k]) + + +@pytest.mark.parametrize("stochastic", [True, False]) +def test_forward_ensemble(stochastic: bool): + torch.manual_seed(0) + model = _make_small_v1p5(stochastic=stochastic) + # 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) + n = 3 + batch = tile_batch(batch, n) + b = next(iter(batch.surf_vars.values())).shape[0] + + with torch.inference_mode(): + pred = model.forward(batch, lead_times=torch.full((b,), 6.0)) + members = split_batch(pred, n) + + if stochastic: + # Every member receives independent noise, so members must differ. + for member1, member2 in itertools.combinations(members, 2): + assert (member1.surf_vars["2t"] - member2.surf_vars["2t"]).abs().mean() > 1e-2 + else: + # Check that all are equal. + for m in range(1, n): + torch.testing.assert_close( + members[0].surf_vars["2t"], members[m].surf_vars["2t"], atol=1e-3, rtol=1e-3 + ) + + +def test_rollout_ensemble(): + num_ensemble_members = 3 + torch.manual_seed(0) + model = _make_small_v1p5(stochastic=True) + _unzero_adaptive_layer_norms(model) # Otherwise, the noise has no effect on the output. + model.eval() + batch = _make_batch(surf_vars=tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF)) + steps = 2 + + with torch.inference_mode(): + preds = list(rollout_ensemble(model, batch, steps, num_ensemble_members)) + + assert len(preds) == steps + for members in preds: + assert len(members) == num_ensemble_members + for member in members: + for v in member.surf_vars.values(): + assert v.shape[0] == BATCH + # Every member receives independent noise, so members must differ. + for member1, member2 in itertools.combinations(members, 2): + assert (member1.surf_vars["2t"] - member2.surf_vars["2t"]).abs().mean() > 1e-2 + + +def test_rollout_ensemble_num_ensemble_members_one_raises(): + num_ensemble_members = 1 + 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) + _ = next(iter(batch.surf_vars.values())).shape[0] + + with pytest.raises(ValueError), torch.inference_mode(): + _ = list(rollout_ensemble(model, batch, steps=2, num_ensemble_members=num_ensemble_members))