File size: 2,986 Bytes
1ca0208 | 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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | """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)
|