File size: 2,851 Bytes
41dd7e0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | """Training pairs and evaluation inputs of a dataset package (``trajectories_fine.npz``, ``split.json``,
``normalization.npz``).
* ``get_dataloaders(cfg, seed)`` -> ``(train_loader, val_loader)``: segment pairs ``(x_tk, x_tk1, t_k, t_k1)`` for
Cascaded SI (``method.data_mode: fine_segment_pairs``) or endpoint pairs ``(x_0, x_T)`` for Direct SI and EDM
(``fine_endpoint_pairs``), from the pair manifest named by ``method.fine_slice_manifest``.
* ``get_split_set(cfg, split, seed, x0_bank)`` -> ``(x0_per_gid, truth_clouds, gids)`` at the forecast horizon
``method.fine_horizon_idx``, the inputs of sampling and of validation.
"""
from __future__ import annotations
from typing import Tuple
import torch
from torch.utils.data import DataLoader
from .fine import (EndpointPairDataset, SegmentPairDataset, ValEndpointPairDataset, # noqa: F401
ValSegmentPairDataset, get_fine_split_set)
DATA_MODES = ("fine_segment_pairs", "fine_endpoint_pairs")
def get_dataloaders(cfg: dict, seed: int = 0) -> Tuple[DataLoader, DataLoader]:
packaged_dir = cfg["data"]["packaged_dataset_path"]
batch_size = int(cfg["training"]["batch_size"])
method_cfg = cfg.get("method", {})
data_mode = method_cfg.get("data_mode")
num_workers = int(cfg["training"].get("num_workers", 0))
if data_mode == "fine_segment_pairs":
train_ds = SegmentPairDataset(packaged_dir, method_cfg["fine_slice_manifest"], seed=seed,
expected_t_values=method_cfg.get("sampling", {}).get("t_values"))
val_ds = ValSegmentPairDataset(packaged_dir, method_cfg["fine_slice_manifest"],
n_members=int(method_cfg.get("fine_val_members", 20)))
elif data_mode == "fine_endpoint_pairs":
train_ds = EndpointPairDataset(packaged_dir, method_cfg["fine_slice_manifest"], seed=seed,
expected_t_values=method_cfg.get("sampling", {}).get("t_values"))
val_ds = ValEndpointPairDataset(packaged_dir, method_cfg["fine_slice_manifest"],
n_members=int(method_cfg.get("fine_val_members", 20)))
else:
raise ValueError(f"method.data_mode must be one of {DATA_MODES}, got {data_mode!r}")
g = torch.Generator()
g.manual_seed(seed)
train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers, generator=g,
pin_memory=num_workers > 0)
val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers,
pin_memory=num_workers > 0)
return train_loader, val_loader
def get_split_set(cfg: dict, split_name: str, seed: int = 0, x0_bank: str | None = None):
return get_fine_split_set(cfg, split_name, seed=seed, x0_bank=x0_bank)
|