| """Data loading and validation for the local SEEDS data contract.""" |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
| from typing import Dict, Iterable, Optional |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import DataLoader, Dataset, DistributedSampler |
|
|
|
|
| class SEEDSDataset(Dataset): |
| def __init__(self, path: str | Path, channels: int, faces: int, height: int, width: int, seed_count: int) -> None: |
| self.path = Path(path) |
| if not self.path.exists(): |
| raise FileNotFoundError(f"dataset does not exist: {self.path}") |
| archive = np.load(self.path, allow_pickle=False) |
| self.seeds = np.asarray(archive["seeds"], dtype=np.float32) |
| self.targets = np.asarray(archive["targets"], dtype=np.float32) |
| self.climate = np.asarray(archive["climate"], dtype=np.float32) |
| expected = (channels, faces, height, width) |
| if self.seeds.ndim != 6 or tuple(self.seeds.shape[2:]) != expected or self.seeds.shape[1] != seed_count: |
| raise ValueError(f"seeds in {self.path} must have shape [N, {seed_count}, {expected}]") |
| if self.targets.ndim != 5 or tuple(self.targets.shape[1:]) != expected: |
| raise ValueError(f"targets in {self.path} must have shape [N, {expected}]") |
| if self.climate.ndim != 5 or tuple(self.climate.shape[1:]) != expected: |
| raise ValueError(f"climate in {self.path} must have shape [N, {expected}]") |
| if len(self.seeds) != len(self.targets) or len(self.targets) != len(self.climate): |
| raise ValueError("seeds, targets and climate must have the same sample count") |
| if not np.isfinite(self.seeds).all() or not np.isfinite(self.targets).all() or not np.isfinite(self.climate).all(): |
| raise ValueError(f"dataset contains NaN or Inf: {self.path}") |
|
|
| def __len__(self) -> int: |
| return len(self.targets) |
|
|
| def __getitem__(self, index: int) -> Dict[str, torch.Tensor]: |
| return { |
| "seeds": torch.from_numpy(self.seeds[index]), |
| "targets": torch.from_numpy(self.targets[index]), |
| "climate": torch.from_numpy(self.climate[index]), |
| } |
|
|
|
|
| def build_dataloader( |
| path: str | Path, |
| channels: int, |
| faces: int, |
| height: int, |
| width: int, |
| seed_count: int, |
| batch_size: int, |
| shuffle: bool, |
| num_workers: int = 0, |
| max_batches: Optional[int] = None, |
| sampler=None, |
| dataset: Optional[SEEDSDataset] = None, |
| ) -> Iterable[Dict[str, torch.Tensor]]: |
| if dataset is None: |
| dataset = SEEDSDataset(path, channels, faces, height, width, seed_count) |
| loader = DataLoader( |
| dataset, |
| batch_size=batch_size, |
| shuffle=shuffle if sampler is None else False, |
| sampler=sampler, |
| num_workers=num_workers, |
| pin_memory=torch.cuda.is_available(), |
| ) |
| if max_batches is None: |
| return loader |
| return (batch for batch_index, batch in enumerate(loader) if batch_index < max_batches) |
|
|
