diff --git a/AGENTS.md b/AGENTS.md index db7f5089..b5b95cbb 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -16,8 +16,8 @@ model implements a ConvNeXt U-Net neural network architecture. We have made sign support quarter degree (0.25 x 0.25 lat/lng) data emulation. The samudra-multi model has an encoder, processor, and decoder structure and aims to emulate ocean physics by first translating data from a physical space to a latent space. The samudra-multi model supports training on multiple scales of data all at once (e.g. one -degree, half degree and quarter degree), either on a "mix" or "match" schedule (i.e. the cross product of each scale for -input and label, or one input/label scale at a time per batch). +degree, half degree and quarter degree) by configuring multiple data sources. Each configured source is trained against +labels from the same source/resolution. ## Data @@ -307,4 +307,4 @@ notebooks/ # Analysis and preprocessing notebooks 5. **Cloud Training**: Supports SkyPilot for remote job execution on AWS & Lambda Labs 6. **Noisy Failure**: Do not swallow errors. If something goes wrong, let it fail loudly. 7. **Avoid Hacks**: Don't accommodate bad designs by adding more cruft -- refactor separately first then make the nice change. -8. **Multi-Scale Support**: samudra-multi supports training on multiple data resolutions simultaneously with "mix" or "match" scheduling +8. **Multi-Scale Support**: samudra-multi supports training on multiple data resolutions simultaneously by configuring multiple matched data sources diff --git a/configs/samudra_multi_om4/train_multiscale.yaml b/configs/samudra_multi_om4/train_multiscale.yaml index e94615f6..151109c6 100644 --- a/configs/samudra_multi_om4/train_multiscale.yaml +++ b/configs/samudra_multi_om4/train_multiscale.yaml @@ -48,7 +48,6 @@ experiment: wandb: mode: disabled project: default - train_schedule: match data: dataset: type: om4 diff --git a/src/samudra/config.py b/src/samudra/config.py index 17365b9f..968a5e80 100644 --- a/src/samudra/config.py +++ b/src/samudra/config.py @@ -780,7 +780,7 @@ def build( corrector = None if len(srcs) != 1: raise ValueError( - 'Samudra only supports training at a single scale! Please set `training_schedule="standard"`.' + "Samudra only supports training at a single scale! Please configure exactly one data source." ) src = srcs[0] if self.corrector is not None: @@ -997,9 +997,6 @@ class DistributedConfig(BaseConfig): dist_backend: str | None = None -TrainSchedule = Literal["standard", "match", "mix"] - - class ExperimentConfig(BaseConfig): name: str = "cm4_samudra" rand_seed: int = 1 @@ -1007,8 +1004,6 @@ class ExperimentConfig(BaseConfig): # we require this to be set by the user but have optional here # so we can leave it out of config files data_root: Location | None = None - # Define multi-scale dataloader example schedule. Default: single scale. - train_schedule: TrainSchedule = "standard" wandb: WandBConfig @cached_property diff --git a/src/samudra/models/samudra_multi.py b/src/samudra/models/samudra_multi.py index dafab1d3..dbcca07e 100644 --- a/src/samudra/models/samudra_multi.py +++ b/src/samudra/models/samudra_multi.py @@ -104,8 +104,6 @@ def forward_once( fts = self.encoder(fts, ctx.input_resolution_cpu) fts = self.processor(fts) - # TODO(alxmrs): When the output resolution differs from the input (i.e. in a "mix" schedule), we cannot use - # residual predictions (`self.pred_residuals` must be `False`). fts = self.decoder(fts, ctx.output_resolution_cpu) # Convert back to float32 diff --git a/src/samudra/train.py b/src/samudra/train.py index 97533d4b..71b4c13a 100644 --- a/src/samudra/train.py +++ b/src/samudra/train.py @@ -4,7 +4,6 @@ import contextlib import datetime -import itertools import logging import multiprocessing import os @@ -12,10 +11,9 @@ import time import warnings from collections import OrderedDict -from collections.abc import Iterable from multiprocessing.context import BaseContext from pathlib import Path -from typing import Any, assert_never +from typing import Any import dask import torch @@ -36,7 +34,7 @@ get_variable_loss_dict, ) from samudra.backend import init_train_backend -from samudra.config import TrainConfig, TrainSchedule, build_loss_fn +from samudra.config import TrainConfig, build_loss_fn from samudra.constants import ( MAX_TRAIN_MODEL_STEPS_FORWARD, BoundaryVarNames, @@ -59,7 +57,7 @@ train_batch, validate_batch, ) -from samudra.utils.data import DataSource, Normalize, get_inference_steps +from samudra.utils.data import Normalize, get_inference_steps from samudra.utils.device import using_gpu from samudra.utils.distributed import ( all_reduce_mean, @@ -150,16 +148,6 @@ def __init__(self, cfg: TrainConfig) -> None: self.data_container = cfg.data.build( data_root=cfg.experiment.resolved_data_root, ) - self.train_schedule: TrainSchedule = cfg.experiment.train_schedule - if self.train_schedule == "mix" and cfg.model.pred_residuals: - raise ValueError( - "Residual predictions on a mixed multiscale training schedule is not currently supported." - ) - if self.train_schedule == "mix" and any(step > 1 for step in cfg.steps): - raise ValueError( - "Step predictions on a mixed multiscale training schedule is not currently supported." - ) - data_num_workers = cfg.data.loading.num_pytorch_workers() persistent_workers = cfg.data.loading.persistent_pytorch_workers() @@ -681,8 +669,8 @@ def validate_one_epoch(self, epoch): and self.validation_images_enabled ) - if self.train_schedule == "standard": - # The standard val aggregator only supports a single scale. + if len(self.data_container.sources) == 1: + # The standard validation aggregator only supports a single scale. val_aggregator = Aggregator.get_validation_aggregator( self.primary_src.metadata, self.hist, @@ -797,23 +785,12 @@ def init_data_loaders(self, cur_step: int) -> None: Args: cur_step: Current training step size """ - scales = self.data_container.sources - match self.train_schedule: - case "standard": - srcs: Iterable[tuple[DataSource, DataSource | None]] = [ - (scales[0], None) - ] - case "match": - srcs = [(s, s) for s in scales] - case "mix": - srcs = list(itertools.product(scales, repeat=2)) # type: ignore - case _: - assert_never(self.train_schedule) + srcs = self.data_container.sources train_datasets = [ TorchTrainDataset( src=src.slice(self.train_time), - dst=dst.slice(self.train_time) if dst else None, + dst=None, prognostic_var_names=self.prognostic_var_names, boundary_var_names=self.boundary_var_names, hist=self.hist, @@ -824,13 +801,13 @@ def init_data_loaders(self, cur_step: int) -> None: concurrent_compute_=self.concurrent_compute, ) for stride in self.data_stride - for src, dst in srcs + for src in srcs ] val_datasets = [ TorchTrainDataset( src=src.slice(self.val_time), - dst=dst.slice(self.val_time) if dst else None, + dst=None, prognostic_var_names=self.prognostic_var_names, boundary_var_names=self.boundary_var_names, hist=self.hist, @@ -841,7 +818,7 @@ def init_data_loaders(self, cur_step: int) -> None: concurrent_compute_=self.concurrent_compute, ) for stride in self.data_stride - for src, dst in srcs + for src in srcs ] # Create datasets @@ -872,7 +849,7 @@ def init_data_loaders(self, cur_step: int) -> None: ) # Create batch samplers - branch on distributed vs non-distributed - # Group by input AND label resolution to handle all training schedules + # Group by resolution so batches stay homogeneous across configured sources. def group_key(ds): return tuple(prog.grid_size for prog in ds.prognostic_srcs) diff --git a/tests/conftest.py b/tests/conftest.py index 0f165180..731b8617 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -16,7 +16,7 @@ from numpy.typing import ArrayLike, NDArray import samudra.constants as c -from samudra.config import JulianDate, TrainBackendConfig, TrainConfig, TrainSchedule +from samudra.config import JulianDate, TrainBackendConfig, TrainConfig from samudra.train import Trainer from samudra.utils.data import DataSource, Masks, _is_compact, compact_dataset from samudra.utils.multiton import MultitonScope @@ -314,11 +314,6 @@ def history(request: pytest.FixtureRequest) -> int: return request.param -@pytest.fixture(scope="session", params=["standard", "match", "mix"]) -def schedule(request: pytest.FixtureRequest) -> TrainSchedule: - return request.param - - # Run a test for both CPU and GPU, and allows selecting or skipping CUDA tests. @pytest.fixture( params=["cpu", pytest.param("cuda", marks=pytest.mark.cuda)], scope="session" diff --git a/tests/test_datasets.py b/tests/test_datasets.py index c501a36f..c0e29062 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -8,8 +8,7 @@ import dataclasses import datetime import itertools -from collections.abc import Generator, Iterable -from typing import assert_never +from collections.abc import Generator import cftime import numpy as np @@ -22,7 +21,7 @@ from numpy.typing import NDArray from torch.utils.data import ConcatDataset, DataLoader -from samudra.config import TimeConfig, TrainConfig, TrainSchedule +from samudra.config import TimeConfig, TrainConfig from samudra.constants import LoaderVersion from samudra.datasets import ( InferenceDataset, @@ -92,7 +91,7 @@ def make_loader( time_config: TimeConfig | None = None, drop_last: bool = True, version: LoaderVersion | None = None, - schedule: TrainSchedule = "standard", + multiscale: bool = False, shuffle: bool = True, ) -> Generator[DataLoader | TrainDataLoader, None, None]: if time_config is None: @@ -116,26 +115,16 @@ def make_loader( pytest.skip(f"{version} does not support compact data.") with MultitonScope(): - match schedule: - case "standard": - srcs: Iterable[tuple[DataSource, DataSource | None]] = [(src, None)] - case "match": - coarsened_src = coarsen_source(src, prognostic, boundary) - scales = [src, coarsened_src] - srcs = [(s, s) for s in scales] - case "mix": - coarsened_src = coarsen_source(src, prognostic, boundary) - scales = [src, coarsened_src] - srcs = list(itertools.product(scales, repeat=2)) # type: ignore - case _: - assert_never(schedule) + srcs = [src] + if multiscale: + srcs.append(coarsen_source(src, prognostic, boundary)) match version: case LoaderVersion.OM4_TORCH: dataset_list = [ TorchTrainDataset( src=src.slice(time_config), - dst=dst.slice(time_config) if dst else None, + dst=None, prognostic_var_names=prognostic, boundary_var_names=boundary, hist=cfg.data.hist, @@ -144,15 +133,14 @@ def make_loader( masked_fill_value=cfg.data.masked_fill_value, stride=stride, ) - for src, dst in srcs + for src in srcs for stride in cfg.data_stride ] data: ConcatDataset = ConcatDataset(dataset_list) collate_fn = collate_raw_train_data - # Group datasets by input AND label resolution, allowing different strides to batch together - # This ensures datasets with same (src, dst) resolution pair but different strides can batch + # Group datasets by resolution, allowing different strides to batch together. batch_sampler = EquivalenceGroupBatchSampler.from_datasets( datasets=dataset_list, group_key=lambda ds: tuple( @@ -192,7 +180,7 @@ def extract_sample_arrays(td: TrainData) -> tuple[np.ndarray, np.ndarray]: def calc_num_samples( - cfg: TrainConfig, time_slice: slice, schedule: TrainSchedule + cfg: TrainConfig, time_slice: slice, source_count: int = 1 ) -> int: primary = cfg.data.sources[0] ds = cfg.experiment.resolved_data_root.resolve(primary.data_location).open() @@ -203,12 +191,7 @@ def calc_num_samples( stride = cfg.data_stride[0] n_samples = data_size - (steps * (cfg.data.hist + 1) * stride) - hist * stride - if schedule == "match": - n_samples *= 2 - if schedule == "mix": - n_samples *= 4 - - return n_samples + return n_samples * source_count def vector_of(max_vec_size: int, min_vec_size=1): @@ -335,9 +318,7 @@ def test_loader__data_shape( ) * num_input_timesteps output_var_dim = len(dataset_spec.prognostic_var_names) * num_input_timesteps - n_samples = calc_num_samples( - train_config, train_config.train_time.time_slice, "standard" - ) + n_samples = calc_num_samples(train_config, train_config.train_time.time_slice) assert len(loader) == n_samples, ( f"Current config {train_config} only supports {n_samples} examples; " f"got {len(loader)}." @@ -366,15 +347,30 @@ def test_loader__data_shape( @pytest.mark.parametrize( "data_source,config_name", [("mock", DEFAULT_CONFIG)], indirect=True ) -def test_loader__data_shape__across_schedules( - train_config: TrainConfig, schedule: TrainSchedule +@pytest.mark.parametrize( + "multiscale,expected_patterns", + [ + (False, {((180, 360), (180, 360))}), + ( + True, + { + ((180, 360), (180, 360)), + ((90, 180), (90, 180)), + }, + ), + ], +) +def test_loader__data_shape__across_source_counts( + train_config: TrainConfig, + multiscale: bool, + expected_patterns: set[tuple[tuple[int, int], tuple[int, int]]], ): history = train_config.data.hist with make_loader( train_config, version=LoaderVersion.OM4_TORCH, - schedule=schedule, + multiscale=multiscale, # Keep grouped-sampler ordering deterministic for resolution coverage. shuffle=False, ) as loader: @@ -389,34 +385,15 @@ def test_loader__data_shape__across_schedules( output_var_dim = len(dataset_spec.prognostic_var_names) * num_input_timesteps n_samples = calc_num_samples( - train_config, train_config.train_time.time_slice, schedule + train_config, + train_config.train_time.time_slice, + source_count=2 if multiscale else 1, ) assert len(loader) == n_samples, ( f"Current config {train_config} only supports {n_samples} examples; " f"got {len(loader)}." ) - match schedule: - case "standard": - expected_patterns = { - ((180, 360), (180, 360)), - } - case "match": - expected_patterns = { - ((180, 360), (180, 360)), - ((90, 180), (90, 180)), - } - case "mix": - # In mix mode with 2 scales, multiplex creates pattern: (0,0), (0,1), (1,0), (1,1) - expected_patterns = { - ((180, 360), (180, 360)), # (0,0): full-res input, full-res label - ((180, 360), (90, 180)), # (0,1): full-res input, half-res label - ((90, 180), (180, 360)), # (1,0): half-res input, full-res label - ((90, 180), (90, 180)), # (1,1): half-res input, half-res label - } - case _: - assert_never(schedule) - min_checked_batches = 10 observed_patterns = set() for batch_idx, sample in enumerate(loader, start=1): @@ -442,7 +419,7 @@ def test_loader__data_shape__across_schedules( break assert observed_patterns == expected_patterns, ( - f"Loader did not produce the expected resolution patterns for {schedule=}. " + f"Loader did not produce the expected resolution patterns for {multiscale=}. " f"Expected: {expected_patterns}, got: {observed_patterns}, " f"missing: {expected_patterns - observed_patterns}, " f"invalid: {observed_patterns - expected_patterns}" @@ -566,13 +543,14 @@ def make_config(src: DataSource): @pytest.mark.parametrize("data_source", ["mock-om4"], indirect=True) -def test_mixed_schedule__has_consistent_collated_batches( - train_config: TrainConfig, schedule: TrainSchedule +@pytest.mark.parametrize("multiscale", [False, True]) +def test_source_count__has_consistent_collated_batches( + train_config: TrainConfig, multiscale: bool ): # Exposes underling consistency issue train_config.batch_size = 4 - with make_loader(train_config, schedule=schedule) as loader: + with make_loader(train_config, multiscale=multiscale) as loader: for _ in itertools.islice(loader, 2): pass