diff --git a/scripts/backtest.py b/scripts/backtest.py index 952d1e2..4f46f3d 100755 --- a/scripts/backtest.py +++ b/scripts/backtest.py @@ -43,6 +43,7 @@ from sat_pred.dataset import SatelliteDataset, TimePeriod from sat_pred.load_model import get_model_from_checkpoints, get_model_from_huggingface from sat_pred.predictions import PREDICTION_DIMS, PREDICTION_VAR_NAME, prediction_coords +from sat_pred.spatial import SpatialGrid @@ -82,6 +83,7 @@ def __init__( self, model: torch.nn.Module, data_config: dict, + spatial_grid: SpatialGrid, model_address: str, device: torch.device, ) -> None: @@ -90,6 +92,8 @@ def __init__( Args: model: The trained torch model data_config: The config of the data the model was trained on + spatial_grid: The grid the model was trained on, for the dataset to select its + inputs onto model_address: Where the model was loaded from. Saved onto the predictions, so a store records which model made it device: The torch device to run the model on @@ -105,6 +109,7 @@ def __init__( self.model = model.to(device).eval() self.device = device self.model_address = model_address + self.spatial_grid = spatial_grid # The shape of the samples the model was trained on. These are unpacked from the data # config rather than the model config so the backtest samples are built exactly the way the @@ -126,8 +131,10 @@ def from_checkpoint(cls, checkpoint_dir_path: str, device: torch.device) -> "MLM checkpoint_dir_path: Path of the checkpoint directory device: The torch device to run the model on """ - model, _, data_config, _ = get_model_from_checkpoints(checkpoint_dir_path, val_best=True) - return cls(model, data_config, checkpoint_dir_path, device) + model, _, data_config, spatial_grid, _ = get_model_from_checkpoints( + checkpoint_dir_path, val_best=True + ) + return cls(model, data_config, spatial_grid, checkpoint_dir_path, device) @classmethod def from_huggingface(cls, repo_id: str, revision: str, device: torch.device) -> "MLModel": @@ -142,8 +149,8 @@ def from_huggingface(cls, repo_id: str, revision: str, device: torch.device) -> revision: The commit hash of the model to download device: The torch device to run the model on """ - model, data_config = get_model_from_huggingface(repo_id, revision) - return cls(model, data_config, f"{repo_id}@{revision}", device) + model, data_config, spatial_grid = get_model_from_huggingface(repo_id, revision) + return cls(model, data_config, spatial_grid, f"{repo_id}@{revision}", device) @torch.no_grad() def __call__(self, X: torch.Tensor) -> torch.Tensor: @@ -267,6 +274,7 @@ def __init__( history_mins: int, sample_freq_mins: int, channels, + spatial_grid: SpatialGrid, ): """A torch Dataset for loading model inputs and the init-time they are a forecast from @@ -279,6 +287,7 @@ def __init__( channels: The channels the model was trained on, in order, and the constants each one is normalised with. Either a mapping of channel name to constants, or the path of a YAML file holding one + spatial_grid: The grid the model was trained on, which the samples are selected onto """ # No target is loaded - the backtest only makes predictions, it does not score them @@ -289,6 +298,7 @@ def __init__( forecast_mins=0, sample_freq_mins=sample_freq_mins, channels=channels, + spatial_grid=spatial_grid, ) # Only forecast on the half hour @@ -520,6 +530,7 @@ def backtest( history_mins=model.history_mins, sample_freq_mins=model.sample_freq_mins, channels=model.channels, + spatial_grid=model.spatial_grid, ) # Without this the run would quietly do nothing and write no store at all, which is easy to diff --git a/scripts/push_checkpoint_to_huggingface.py b/scripts/push_checkpoint_to_huggingface.py index e995bfe..2a4189c 100755 --- a/scripts/push_checkpoint_to_huggingface.py +++ b/scripts/push_checkpoint_to_huggingface.py @@ -23,8 +23,10 @@ MODEL_CARD_NAME, MODEL_CONFIG_NAME, PYTORCH_WEIGHTS_NAME, + SPATIAL_GRID_NAME, ) from sat_pred.load_model import get_model_from_checkpoints +from sat_pred.spatial import SpatialGrid from pathlib import Path @@ -46,6 +48,7 @@ def save_model_to_huggingface( save_directory: str, model_config: dict, data_config: dict, + spatial_grid: SpatialGrid, wandb_repo: str, wandb_id: str, experiment_config_path: str | None = None, @@ -65,6 +68,8 @@ def save_model_to_huggingface( Model configuration specified as a key/value dictionary. data_config: Data configuration the model was trained with, specified as a key/value dictionary. + spatial_grid: + The grid the model was trained on. wandb_repo: Identifier of the repo on wandb. wandb_id: Identifier of the model on wandb. experiment_config_path: Path to the full hydra config of the training run, if it was saved. @@ -90,6 +95,11 @@ def save_model_to_huggingface( with open(save_directory / DATA_CONFIG_NAME, 'w') as outfile: yaml.dump(data_config, outfile, default_flow_style=False) + # Save the grid the model was trained on. Unlike the experiment config below this is not + # optional - without it nothing downstream can check that an input covers the area the model + # was trained on + spatial_grid.save(save_directory / SPATIAL_GRID_NAME) + # Save the full config of the training run, if it was saved with the checkpoint if experiment_config_path is not None: shutil.copyfile(experiment_config_path, save_directory / FULL_CONFIG_NAME) @@ -168,9 +178,8 @@ def push_to_huggingface( raise ValueError(f"Could not find wandb run '{wandb_id}' within {wandb_repo}") # Load the model - model, model_config, data_config, experiment_config_path = get_model_from_checkpoints( - checkpoint_dir_path, - val_best=val_best + model, model_config, data_config, spatial_grid, experiment_config_path = ( + get_model_from_checkpoints(checkpoint_dir_path, val_best=val_best) ) # Push to hub @@ -185,6 +194,7 @@ def push_to_huggingface( save_directory=model_output_dir, model_config=model_config, data_config=data_config, + spatial_grid=spatial_grid, experiment_config_path=experiment_config_path, wandb_repo=wandb_repo, wandb_id=wandb_id, diff --git a/src/sat_pred/constants.py b/src/sat_pred/constants.py index ad2d2c4..14f3ecd 100644 --- a/src/sat_pred/constants.py +++ b/src/sat_pred/constants.py @@ -10,3 +10,4 @@ FULL_CONFIG_NAME = "full_experiment_config.yaml" PYTORCH_WEIGHTS_NAME = "model.safetensors" MODEL_CARD_NAME = "README.md" +SPATIAL_GRID_NAME = "spatial_grid.npz" diff --git a/src/sat_pred/dataset.py b/src/sat_pred/dataset.py index 155a2e3..39ba188 100644 --- a/src/sat_pred/dataset.py +++ b/src/sat_pred/dataset.py @@ -21,6 +21,7 @@ from ocf_data_sampler.select.find_contiguous_time_periods import find_contiguous_t0_periods from ocf_data_sampler.torch_datasets.pvnet_dataset import PickleCacheMixin from sat_pred.channels import ChannelConfigInput, parse_channel_config +from sat_pred.spatial import SpatialGrid from sat_pred.xr_tensorstore import open_zarr_paths @@ -137,6 +138,7 @@ def __init__( forecast_mins: int, sample_freq_mins: int, channels: ChannelConfigInput, + spatial_grid: SpatialGrid | None = None, preshuffle: bool = False, seed: int | None = None, ): @@ -154,6 +156,8 @@ def __init__( the constants each one is clipped and z-scored with. Either a mapping of channel name to constants, or the path of a YAML file holding one - see `configs/datamodule/channels/` + spatial_grid: The grid a model was trained on, if the samples should be selected onto + it. `None` takes whatever area the store holds preshuffle (bool): Whether to shuffle the data - useful for validation. Defaults to False. seed (int | None): Seed used to shuffle the data if `preshuffle` is True. Set this to @@ -172,6 +176,11 @@ def __init__( # the store holds them in, so a model always reads the same channel at the same index da = da.sel(channel=self.channel_config.names) + # Selecting the model's own grid out of the store, so an archive covering a different area + # is cropped to what the model was trained on, or refused if it cannot cover it + if spatial_grid is not None: + da = spatial_grid.select(da) + # Convert the satellite data to the given time frequency by selection mask = np.mod(da.time_utc.dt.minute, sample_freq_mins) == 0 da = da.sel(time_utc=mask) @@ -198,6 +207,11 @@ def __init__( self.da = da self.t0_times = t0_times + @property + def spatial_grid(self) -> SpatialGrid: + """The grid the samples are taken from""" + return SpatialGrid.from_dataarray(self.da) + def __setstate__(self, state: dict) -> None: """Rebuild a dataset sent to a dataloader worker @@ -343,6 +357,11 @@ def __init__( # Parsed once here so a bad config fails before any data is loaded self.channel_config = parse_channel_config(channels) + @property + def spatial_grid(self) -> SpatialGrid: + """The grid the training samples are taken from""" + return self._get_dataset("train").spatial_grid + def _make_dataset( self, time_periods: list[TimePeriod], preshuffle: bool = False ) -> SatelliteDataset: diff --git a/src/sat_pred/load_model.py b/src/sat_pred/load_model.py index 44b29fc..aaf7004 100755 --- a/src/sat_pred/load_model.py +++ b/src/sat_pred/load_model.py @@ -10,7 +10,6 @@ device and call `.eval()` themselves. """ -import os from glob import glob import hydra @@ -24,7 +23,9 @@ FULL_CONFIG_NAME, MODEL_CONFIG_NAME, PYTORCH_WEIGHTS_NAME, + SPATIAL_GRID_NAME, ) +from sat_pred.spatial import SpatialGrid def _read_yaml_config(path: str) -> dict: @@ -60,7 +61,7 @@ def get_checkpoint_path(checkpoint_dir_path: str, val_best: bool = True) -> str: def get_model_from_checkpoints( checkpoint_dir_path: str, val_best: bool = True, -) -> tuple[torch.nn.Module, dict, dict, str]: +) -> tuple[torch.nn.Module, dict, dict, SpatialGrid, str]: """Load a model from its checkpoint directory Args: @@ -73,6 +74,7 @@ def get_model_from_checkpoints( model: The trained torch model, with the lightning wrapper discarded model_config: The config of the torch model data_config: The config of the data the model was trained on + spatial_grid: The grid the model was trained on experiment_config_path: Path to the full hydra config of the training run. This is None for models trained before these configs were saved """ @@ -95,15 +97,17 @@ def get_model_from_checkpoints( data_config = _read_yaml_config(f"{checkpoint_dir_path}/{DATA_CONFIG_NAME}") + spatial_grid = SpatialGrid.load(f"{checkpoint_dir_path}/{SPATIAL_GRID_NAME}") + experiment_config_path = f"{checkpoint_dir_path}/{FULL_CONFIG_NAME}" - return model, model_config, data_config, experiment_config_path + return model, model_config, data_config, spatial_grid, experiment_config_path def get_model_from_huggingface( repo_id: str, revision: str, -) -> tuple[torch.nn.Module, dict]: +) -> tuple[torch.nn.Module, dict, SpatialGrid]: """Load a model from a huggingface model repo The repo holds what `scripts/push_checkpoint_to_huggingface.py` pushed. Note that its @@ -125,6 +129,8 @@ def get_model_from_huggingface( data_config: The config of the data the model was trained on. This gives the history, forecast horizon, time resolution and channels the model expects, so the inputs built for it always match what it was trained to read + spatial_grid: The grid the model was trained on, for the caller to select its inputs + onto """ download_dir = snapshot_download(repo_id=repo_id, revision=revision) @@ -141,4 +147,6 @@ def get_model_from_huggingface( data_config = _read_yaml_config(f"{download_dir}/{DATA_CONFIG_NAME}") - return model, data_config + spatial_grid = SpatialGrid.load(f"{download_dir}/{SPATIAL_GRID_NAME}") + + return model, data_config, spatial_grid diff --git a/src/sat_pred/spatial.py b/src/sat_pred/spatial.py new file mode 100644 index 0000000..86693f2 --- /dev/null +++ b/src/sat_pred/spatial.py @@ -0,0 +1,84 @@ +"""The spatial grid a model was trained on, and selecting model inputs onto it + +The training data is a crop of the full satellite disc - for the UK models, 614 x 372 pixels at +about 3km. Nothing about that crop is recorded in the model config, and the models are fully +convolutional, so one handed a different crop at inference time runs happily and predicts plausible +values over the wrong geography. The grid is therefore recorded alongside the model when training +starts, and the inputs are selected onto it before they reach the model: an archive covering more +than the model needs is cropped to it, and one which cannot cover it is refused. + +The coordinates are recorded as full arrays rather than as an extent and a size. The grid is +regular, so an extent would determine it, but rebuilding it with `linspace` would not reproduce the +original floats bit for bit, and the selection is an exact lookup. +""" + +from pathlib import Path + +import numpy as np +import xarray as xr +from numpy.typing import NDArray + + +class SpatialGrid: + """The spatial grid a model was trained on""" + + def __init__( + self, + x_geostationary: NDArray[np.float64], + y_geostationary: NDArray[np.float64], + ) -> None: + """The spatial grid a model was trained on + + Both coordinates must be as `open_sat_data` leaves them. It runs + `make_spatial_coords_increasing`, which reverses `x_geostationary` - the stores hold it + descending - so a grid recorded on one side of that transform would never match one + recorded on the other. + + Args: + x_geostationary: The x coordinate of every pixel column + y_geostationary: The y coordinate of every pixel row + """ + self.x_geostationary = x_geostationary + self.y_geostationary = y_geostationary + + @classmethod + def from_dataarray(cls, da: xr.DataArray) -> "SpatialGrid": + """The grid some satellite data is on + + Args: + da: Satellite data as `open_sat_data` returns it + """ + return cls(da.x_geostationary.values, da.y_geostationary.values) + + @classmethod + def load(cls, path: str | Path) -> "SpatialGrid": + """Load a grid saved by `save` + + Args: + path: Path of the file to load + """ + with np.load(path) as file: + return cls(file["x_geostationary"], file["y_geostationary"]) + + def save(self, path: str | Path) -> None: + """Save the grid so that `load` can read it back + + Numpy's own container rather than one of the YAML configs saved beside it, because two + float64 coordinate arrays round-trip through it exactly, without the check having to depend + on how a text format happens to render a float. + + Args: + path: Path of the file to write. Numpy appends `.npz` if it is not already there + """ + np.savez(path, x_geostationary=self.x_geostationary, y_geostationary=self.y_geostationary) + + def select(self, da: xr.DataArray) -> xr.DataArray: + """Select the pixels of this grid out of some satellite data + + Data covering more than the grid is cropped to it. Data which does not cover all of it + raises, because a model cannot be run over an area it was never trained on. + + Args: + da: Satellite data as `open_sat_data` returns it + """ + return da.sel(x_geostationary=self.x_geostationary, y_geostationary=self.y_geostationary) diff --git a/src/sat_pred/train.py b/src/sat_pred/train.py index ea16b55..25f6ad2 100755 --- a/src/sat_pred/train.py +++ b/src/sat_pred/train.py @@ -2,6 +2,7 @@ import os import hydra +from pathlib import Path from lightning.pytorch import ( Callback, LightningDataModule, @@ -18,7 +19,12 @@ import rich.tree from lightning.pytorch.utilities import rank_zero_only -from sat_pred.constants import DATA_CONFIG_NAME, FULL_CONFIG_NAME, MODEL_CONFIG_NAME +from sat_pred.constants import ( + DATA_CONFIG_NAME, + FULL_CONFIG_NAME, + MODEL_CONFIG_NAME, + SPATIAL_GRID_NAME, +) from sat_pred.load_model import get_checkpoint_path, get_model_from_checkpoints from sat_pred.loss import LossFunction @@ -62,10 +68,11 @@ def print_config( """ style = "dim" - tree = rich.tree.Tree("CONFIG", style=style, guide_style=style) + style_kwargs = {"style": style, "guide_style": style} + tree = rich.tree.Tree("CONFIG", **style_kwargs) for field in fields: - branch = tree.add(field, style=style, guide_style=style) + branch = tree.add(field, **style_kwargs) config_section = config.get(field) branch_content = str(config_section) @@ -77,6 +84,19 @@ def print_config( rich.print(tree) +def get_next_instance(sequence, class_type) -> object | None: + """Get the next instance of a class in a sequence, or None if there is none. + + Args: + sequence: A sequence of objects to search through. + class_type: The class type to look for. + """ + for item in sequence: + if isinstance(item, class_type): + return item + return None + + @hydra.main(config_path="../../configs/", config_name="config.yaml", version_base="1.2") def train(config: DictConfig): """Train the model using parameters in the supplied config files. @@ -98,7 +118,7 @@ def train(config: DictConfig): val_best = config.model.model.val_best # Load the model from the checkpoint - torch_model, model_config, _, _ = get_model_from_checkpoints( + torch_model, model_config, *_ = get_model_from_checkpoints( checkpoint_dir, val_best=val_best ) @@ -135,16 +155,14 @@ def train(config: DictConfig): if "_target_" in cb_conf: callbacks.append(hydra.utils.instantiate(cb_conf)) + # Instantiate the datamodule + datamodule: LightningDataModule = hydra.utils.instantiate(config.datamodule, _convert_='all') + # Align the wandb id with the checkpoint path # - only works if wandb logger and model checkpoint used - use_wandb_logger = False - for logger in loggers: - if isinstance(logger, WandbLogger): - use_wandb_logger = True - wandb_logger = logger - break - - if use_wandb_logger: + wandb_logger = get_next_instance(loggers, WandbLogger) + + if wandb_logger is not None: # Calling the .experiment property initialises the logger wandb_run = wandb_logger.experiment @@ -156,34 +174,34 @@ def train(config: DictConfig): wandb_run.define_metric("trainer/global_step") wandb_run.define_metric("*", step_metric="trainer/global_step") - for callback in callbacks: - if isinstance(callback, ModelCheckpoint): - # skip for non-rank-0 processes: - # see https://github.com/Lightning-AI/pytorch-lightning/issues/13166#issuecomment-1139765549 - if wandb_logger.version is None: - break + checkpoint_callback = get_next_instance(callbacks, ModelCheckpoint) - callback.dirpath = "/".join( - callback.dirpath.split("/")[:-1] + [wandb_logger.version] - ) - # Also save model config to this path - os.makedirs(callback.dirpath, exist_ok=True) - OmegaConf.save(config.model, f"{callback.dirpath}/{MODEL_CONFIG_NAME}") + # A version of None means a non-rank-0 process, which must not write any of this: + # see https://github.com/Lightning-AI/pytorch-lightning/issues/13166#issuecomment-1139765549 + if checkpoint_callback is not None and wandb_logger.version is not None: - # Similarly save the data config - OmegaConf.save(config.datamodule, f"{callback.dirpath}/{DATA_CONFIG_NAME}") + dirpath = str(Path(checkpoint_callback.dirpath).with_name(wandb_logger.version)) + os.makedirs(dirpath, exist_ok=True) - # Save the full resolved hydra config to this path and upload it to wandb - full_config_path = f"{callback.dirpath}/{FULL_CONFIG_NAME}" - OmegaConf.save(config, full_config_path, resolve=True) - wandb_logger.experiment.save(full_config_path, base_path=callback.dirpath) + # Set the checkpoint callback so it writes to this path + checkpoint_callback.dirpath = dirpath - break + # Also save model config to this path + OmegaConf.save(config.model, f"{dirpath}/{MODEL_CONFIG_NAME}") - # Instantiate the datamodule - datamodule: LightningDataModule = hydra.utils.instantiate(config.datamodule, _convert_='all') - - datamodule.zarr_path = list(datamodule.zarr_path) + # Similarly save the data config + OmegaConf.save(config.datamodule, f"{dirpath}/{DATA_CONFIG_NAME}") + + # And the grid the model is about to be trained on, which the data config does not + # record. Saved now rather than when the model is pushed, because the zarr paths could + # have gone stale by then. This builds the train dataset, so the valid-t0 search + # happens here instead of in `fit` - it is cached, so `fit` reuses it + datamodule.spatial_grid.save(f"{dirpath}/{SPATIAL_GRID_NAME}") + + # Save the full resolved hydra config to this path and upload it to wandb + full_config_path = f"{dirpath}/{FULL_CONFIG_NAME}" + OmegaConf.save(config, full_config_path, resolve=True) + wandb_logger.experiment.save(full_config_path, base_path=dirpath) trainer: Trainer = hydra.utils.instantiate( config.trainer, diff --git a/tests/conftest.py b/tests/conftest.py index 19a780d..6bd8383 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -22,7 +22,9 @@ FULL_CONFIG_NAME, MODEL_CONFIG_NAME, PYTORCH_WEIGHTS_NAME, + SPATIAL_GRID_NAME, ) +from sat_pred.spatial import SpatialGrid # The channel config the training configs use CHANNELS_PATH = Path(__file__).parents[1] / "configs/datamodule/channels/seviri_rss.yaml" @@ -51,6 +53,13 @@ "IR_134", "VIS006", "VIS008", "WV_062", "WV_073", ] +# The grid of the synthetic stores, with `x_geostationary` ascending as `open_sat_data` leaves it +# rather than descending as the stores hold it +SPATIAL_GRID = SpatialGrid( + x_geostationary=np.sort(np.arange(IMAGE_SIZE, dtype=np.float64) * -3000), + y_geostationary=np.arange(IMAGE_SIZE, dtype=np.float64) * 3000, +) + AREA_STRING = json.dumps( { "msg_seviri_rss_3km": { @@ -224,6 +233,8 @@ def checkpoint_dir(tmp_path, training_module_config) -> str: yaml.dump({"seed": 12345, "model": training_module_config}) ) + SPATIAL_GRID.save(tmp_path / SPATIAL_GRID_NAME) + torch.save({"state_dict": training_module.state_dict()}, tmp_path / "epoch=1-step=10.ckpt") return str(tmp_path) @@ -246,4 +257,6 @@ def huggingface_dir(tmp_path, tiny_model_config) -> str: save_model_as_safetensor(model, str(tmp_path / PYTORCH_WEIGHTS_NAME)) + SPATIAL_GRID.save(tmp_path / SPATIAL_GRID_NAME) + return str(tmp_path) diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 8634efd..401e315 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -7,6 +7,7 @@ import pytest from sat_pred.dataset import SatelliteDataModule, SatelliteDataset, find_valid_t0_times +from sat_pred.spatial import SpatialGrid from tests.conftest import ( CHANNELS_PATH, DATA_FREQ_MINS, @@ -138,6 +139,28 @@ def test_a_channel_the_data_does_not_have_is_an_error(sat_zarr_path, channel_con ) +def test_a_dataset_can_be_selected_onto_a_grid(sat_zarr_path): + """Samples are cropped to the grid given, and the dataset reports the grid it is on + + Recording the grid at training and selecting onto it at inference only line up if these two + are inverses of each other. + """ + full_dataset = make_dataset(sat_zarr_path, [PERIOD_1]) + + grid = SpatialGrid( + x_geostationary=full_dataset.da.x_geostationary.values[2:10], + y_geostationary=full_dataset.da.y_geostationary.values[3:12], + ) + + dataset = make_dataset(sat_zarr_path, [PERIOD_1], spatial_grid=grid) + + np.testing.assert_array_equal(dataset.spatial_grid.x_geostationary, grid.x_geostationary) + np.testing.assert_array_equal(dataset.spatial_grid.y_geostationary, grid.y_geostationary) + + X, _ = dataset[0] + assert X.shape == (NUM_CHANNELS, NUM_HISTORY_STEPS, 9, 8) + + def test_samples_are_normalised(sat_zarr_path, channel_config): """Samples come out as per-channel z-scores of the clipped physical values""" dataset = make_dataset(sat_zarr_path, [PERIOD_1]) diff --git a/tests/test_load_model.py b/tests/test_load_model.py index 00de9be..c11328b 100644 --- a/tests/test_load_model.py +++ b/tests/test_load_model.py @@ -11,7 +11,7 @@ from sat_pred.models.simvp_model import SimVP from sat_pred.training_module import TrainingModule -from .conftest import HISTORY_MINS +from .conftest import HISTORY_MINS, IMAGE_SIZE # Stand-ins for the repo and pinned commit a caller would name HF_REPO_ID = "some-org/some-model" @@ -20,8 +20,8 @@ def test_get_model_from_checkpoints(checkpoint_dir, tiny_model_config): """A saved checkpoint round-trips back to a model and its configs""" - model, model_config, data_config, experiment_config_path = get_model_from_checkpoints( - checkpoint_dir + model, model_config, data_config, spatial_grid, experiment_config_path = ( + get_model_from_checkpoints(checkpoint_dir) ) # The lightning wrapper is discarded, leaving the torch model @@ -33,6 +33,7 @@ def test_get_model_from_checkpoints(checkpoint_dir, tiny_model_config): assert model_config == tiny_model_config assert data_config["history_mins"] > 0 + assert len(spatial_grid.x_geostationary) == IMAGE_SIZE assert experiment_config_path == f"{checkpoint_dir}/{FULL_CONFIG_NAME}" @@ -76,8 +77,8 @@ def fake_snapshot_download(repo_id: str, revision: str) -> str: def test_get_model_from_huggingface(downloaded_huggingface_dir): - """A model pushed to huggingface round-trips back to a model and its data config""" - model, data_config = get_model_from_huggingface(HF_REPO_ID, revision=HF_REVISION) + """A model pushed to huggingface round-trips back to a model, data config and grid""" + model, data_config, spatial_grid = get_model_from_huggingface(HF_REPO_ID, HF_REVISION) # The huggingface repo holds the bare model, whereas a checkpoint directory holds the training # module wrapping it - so a loader which conflated the two would fail here @@ -86,6 +87,9 @@ def test_get_model_from_huggingface(downloaded_huggingface_dir): assert data_config["history_mins"] == HISTORY_MINS + assert len(spatial_grid.x_geostationary) == IMAGE_SIZE + assert len(spatial_grid.y_geostationary) == IMAGE_SIZE + # `strict=True` already catches a key mismatch. This catches the weights never being loaded at # all, which leaves the model randomly initialised and forecasting silently wrong weights_path = f"{downloaded_huggingface_dir}/{PYTORCH_WEIGHTS_NAME}"