File size: 2,662 Bytes
275b5a1 | 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 | 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
|