Evaluation pipeline for ocean emulator models.
Runs long autoregressive rollouts and computes metrics against ground-truth
ocean states.
Source code in src/samudra/eval.py
| def __init__(self, cfg: EvalConfig) -> None:
cfg.prepare_output_dirs()
self.device = init_eval_backend(cfg.backend)
# Adjust workers and memory pinning based on device
data_num_workers = cfg.data.loading.num_pytorch_workers()
if not using_gpu():
data_num_workers = 0 # Disable multi-processing on CPU
elif cfg.disk_mode:
data_num_workers = torch.cuda.device_count() * data_num_workers
# Set seeds
set_seed(cfg.experiment.rand_seed)
logger.info("Loading data")
self.data_bundle = cfg.data.build(
cfg.experiment.resolved_data_root,
)
# Getting prognostic and boundary variables
self.data_layout = self.data_bundle.data_layout
self.prognostic_var_names: PrognosticVarNames = (
self.data_layout.prognostic_var_names
)
self.boundary_var_names: BoundaryVarNames = self.data_layout.boundary_var_names
self.levels = self.data_layout.num_prognostic_depth_levels
str_prognostics = ", ".join([i for i in self.prognostic_var_names])
str_boundaries = ", ".join([i for i in self.boundary_var_names])
logger.info(f"Prognostic variables: {str_prognostics}")
logger.info(f"Boundary variables: {str_boundaries}")
logger.info(f"Levels: {self.levels}")
self.N_bound = len(self.boundary_var_names)
self.N_prog = len(self.prognostic_var_names)
self.input_steps = cfg.data.input_steps
self.output_steps = cfg.data.output_steps
self.num_prog_in = self.input_steps * self.N_prog
self.num_boundary_in = self.input_steps * self.N_bound
self.num_in = self.num_prog_in + self.num_boundary_in
self.num_out = self.output_steps * self.N_prog
self.data_layout = self.data_layout.to(self.device)
logger.info(f"Number of inputs (prognostic + boundary): {self.num_in}")
logger.info(f"Number of outputs (prognostic): {self.num_out}")
# Dataloaders
if self.data_bundle.inference_source is None:
raise ValueError(
"Inference time is not configured for the first data source"
)
self.source = self.data_bundle.inference_source
self.metadata = self.source.metadata
self.wet = self.source.masks.prognostic_for_steps(self.output_steps)
self.area_weights: Grid = self.source.spherical_area_weights
self.area_weights = self.area_weights.to(self.device)
self.preprocessor = BatchPreprocessor(
self.source,
prognostic_var_names=self.prognostic_var_names,
boundary_var_names=self.boundary_var_names,
)
# Model
self.model = cfg.model.build(
prog_channels=self.num_prog_in,
boundary_channels=self.num_boundary_in,
out_channels=self.num_out,
input_steps=self.input_steps,
grid_sizes=[source.grid_size for source in self.data_bundle.train_sources],
).to(self.device)
get_model_summary(self.model, None, cfg.debug)
if cfg.ckpt_path is None:
raise ValueError(
"ckpt_path must be set; try --ckpt_path=path/to/checkpoint"
)
self.load_checkpoint(cfg.ckpt_path)
self.network = self.model.__class__.__name__
# Initialize WandB
self.wandb_logger = WandBLogger.init_instance()
self.wandb_logger.configure(
cfg.experiment.wandb.mode == "online", is_main_process()
)
# Set up wandb run
self.wandb_id, self.wandb_name = self.wandb_logger.setup_run(
None, cfg, data_bundle=self.data_bundle, finetune=False
)
# Eval
self.output_dir = cfg.experiment.output_dir
self.debug = cfg.debug
self.num_workers = data_num_workers
self.num_model_steps_forward = cfg.num_model_steps_forward
self.save_zarr = cfg.save_zarr
self.model_path = cfg.ckpt_path
self.normalize_before_mask = cfg.data.normalize_before_mask
self.masked_fill_value = cfg.data.masked_fill_value
self.observations = cfg.observations
self.resolved_data_root = cfg.experiment.resolved_data_root
self.experiment_name = cfg.experiment.name
self.init_inference_store()
|
report_observation_metrics(obs_cfg)
Score the finished rollout against observation products.
Source code in src/samudra/eval.py
| def report_observation_metrics(self, obs_cfg: ObsMetricsConfig) -> None:
"""Score the finished rollout against observation products."""
logger.info("Computing observation metrics")
baselines = {}
if "om4" in obs_cfg.baselines:
# The ground-truth data the eval already staged. The metric table
# is only interpretable next to it, so scoring it costs one extra
# pass over data that is local anyway.
baselines["om4"] = self.source.to_xarray_dataset()
frame, scalars = run_observation_metrics(
obs_cfg,
predictions=open_predictions(self.output_dir),
data_layout=self.data_layout,
data_root=self.resolved_data_root,
model_label=self.experiment_name,
baselines=baselines,
output_dir=self.output_dir,
)
self.wandb_logger.log(
{**scalars, "obs/metrics_table": self.wandb_logger.Table(dataframe=frame)},
step=None,
)
|