| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| import h5py |
| import numpy as np |
| import torch |
| from torch.utils.data import Dataset |
|
|
|
|
| class TwoConditionDataset(Dataset): |
| def __init__(self, path: Path, cache_in_memory: bool = False): |
| self.path = Path(path) |
| self.cache_in_memory = bool(cache_in_memory) |
| with h5py.File(self.path, 'r') as f: |
| self.data_shape = tuple(f['data'].shape) |
| self.label_shape = tuple(f['labels'].shape) |
| self.dtype = f['data'].dtype |
| if self.data_shape != (3000, 6000, 3) or self.label_shape != (3000, 2): |
| raise ValueError(f'unexpected shapes: data={self.data_shape}, labels={self.label_shape}') |
| self._file = None |
| self._data = None |
| self._labels = None |
| self._waveforms = None |
| self._cached_labels = None |
| if self.cache_in_memory: |
| with h5py.File(self.path, 'r') as f: |
| waveforms = np.asarray(f['data'], dtype=np.float32) |
| labels = np.asarray(f['labels'], dtype=np.float32) |
| if not np.isfinite(waveforms).all() or not np.isfinite(labels).all(): |
| raise ValueError('training HDF5 contains non-finite values') |
| self._waveforms = np.ascontiguousarray(waveforms.transpose(0, 2, 1)) |
| self._cached_labels = np.ascontiguousarray(labels) |
| print( |
| f'Cached training data in RAM: waveforms={self._waveforms.nbytes / 2**20:.1f} MiB, ' |
| f'labels={self._cached_labels.nbytes / 2**20:.1f} MiB' |
| ) |
|
|
| def __len__(self): |
| return self.data_shape[0] |
|
|
| def _open(self): |
| if self._file is None: |
| self._file = h5py.File(self.path, 'r') |
| self._data, self._labels = self._file['data'], self._file['labels'] |
|
|
| def __getitem__(self, i): |
| if self._waveforms is not None: |
| x = self._waveforms[i] |
| raw = self._cached_labels[i] |
| else: |
| self._open() |
| x = np.asarray(self._data[i], dtype=np.float32) |
| raw = np.asarray(self._labels[i], dtype=np.float32) |
| if not np.isfinite(x).all() or not np.isfinite(raw).all(): |
| raise ValueError(f'non-finite row {i}') |
| label = raw |
| if not (0 <= label[0] < label[1] <= 6000): |
| raise ValueError(f'invalid label row {i}: {label.tolist()}') |
| waveform = x if self._waveforms is not None else x.transpose(1, 0).copy() |
| return torch.from_numpy(waveform), torch.from_numpy(label) |
|
|
| def close(self): |
| if self._file is not None: |
| self._file.close() |
| self._file = None |
|
|
