diff --git a/ml4ht/data/data_source.py b/ml4ht/data/data_source.py new file mode 100644 index 0000000..df1ad9c --- /dev/null +++ b/ml4ht/data/data_source.py @@ -0,0 +1,200 @@ +from typing import ( + Any, + Dict, + List, + Tuple, + Callable, + Optional, + Sequence, + Iterator, + Mapping, +) + +from torch.utils.data import IterableDataset, Dataset, get_worker_info +import numpy as np +from ml4ht.data.defines import EXCEPTIONS + + +DataIndex = Mapping[str, Any] +Data = Dict[str, np.ndarray] +# Given `DataIndex`, return (input data, output data) +DataSource = Callable[[DataIndex], Tuple[Data, Data]] +TrainingExample = Tuple[ + Dict[str, np.ndarray], + Dict[str, np.ndarray], +] + + +def range_epoch_idx_generator( + true_epoch_len: int, + shuffle: bool, + training_epochs_per_true_epoch: int = 1, +) -> Iterator[Sequence[int]]: + """ + Yields optionally shuffled lists of indices from a range + :param true_epoch_len: how many indices to yield from + :param shuffle: Whether to shuffle the sample ids each epoch + :param training_epochs_per_true_epoch: How many training epochs to break a true epoch into. + Defaults to 1 to yield true epochs. + :yield: Epoch worth of sample ids + """ + assert training_epochs_per_true_epoch <= true_epoch_len + sample_ids = np.arange(true_epoch_len) + while True: + if shuffle: + np.random.shuffle(sample_ids) + epochs = np.array_split(sample_ids, training_epochs_per_true_epoch) + yield from epochs + + +def data_index_epoch_generator( + data_indices: List[DataIndex], + shuffle: bool, + training_epochs_per_true_epoch: int = 1, +) -> Iterator[Sequence[DataIndex]]: + """ + Yields optionally shuffled epochs of `DataIndex`s + :param data_indices: sample ids (e.g. MRNs) for each epoch + :param shuffle: Whether to shuffle the sample ids each epoch + :param training_epochs_per_true_epoch: How many training epochs to break a true epoch into. + Defaults to 1 to yield true epochs. + :yield: Epoch worth of DataIndices + """ + epoch_idx_generator = range_epoch_idx_generator( + true_epoch_len=len(data_indices), + shuffle=shuffle, + training_epochs_per_true_epoch=training_epochs_per_true_epoch, + ) + for indices in epoch_idx_generator: + yield [data_indices[idx] for idx in indices] + + +class TrainingDataset(Dataset): + """ + Indexable of training data from `DataSources` for use in pytorch `DataLoader`. + To be able to skip errors during data loading, use `TrainingIterableDataset` below. + """ + + def __init__( + self, + data_sources: List[DataSource], + data_indices: Sequence[DataIndex], + ): + """ + :param data_sources: Sources of data for inputs and outputs to model + :param data_indices: Indices for each epoch of data + """ + super(TrainingIterableDataset).__init__() + self.data_sources = data_sources + self.data_indices = data_indices + + def __len__(self): + return len(self.data_indices) + + def training_example(self, data_idx: DataIndex) -> Tuple[Data, Data]: + combined_inputs, combined_outputs = {}, {} + for data_source in self.data_sources: + data_in, data_out = data_source(data_idx) + for name, modality in data_in.items(): + if name in combined_inputs: + raise ValueError( + f"Multiple DataSources share the same input modality called {name}", + ) + combined_inputs[name] = modality + for name, modality in data_out.items(): + if name in combined_outputs: + raise ValueError( + f"Multiple DataSources share the same output modality called {name}", + ) + combined_outputs[name] = modality + return combined_inputs, combined_outputs + + def __getitem__(self, item: int) -> Tuple[Data, Data]: + data_idx = self.data_indices[item] + return self.training_example(data_idx) + + +class TrainingIterableDataset(IterableDataset): + """ + Iterable over epochs of training data from `DataSources` for use in pytorch `DataLoader`. + More flexible version of `TrainingDataset` that can skip errors, + but requires defining an Iterator over epochs of `DataIndex`s. + """ + + # keys in error list + error_key = "error" + source_key = "source" + idx_key = "data_idx" + + def __init__( + self, + data_sources: List[DataSource], + epoch_indices_iterator: Iterator[Sequence[DataIndex]], + raise_errors: bool = True, + verbose: bool = True, + ): + """ + :param data_sources: Sources of data for inputs and outputs to model + :param epoch_indices_iterator: Iterator over epochs of `DataIndex`s + For example, see `sample_id_epoch_indices` above. + :param raise_errors: Whether to raise or skip errors during calling of + `DataFetcher`s + """ + super(TrainingIterableDataset).__init__() + self.data_sources = data_sources + self.epoch_indices_iterator = epoch_indices_iterator + self.raise_errors = raise_errors + self.verbose = verbose + self.errors = [] + + def __iter__(self): + """ + One epoch of training data. + """ + epoch_indices = next(self.epoch_indices_iterator) + worker_info = get_worker_info() + if worker_info is not None: # multi-process case + split_indices = np.array_split(epoch_indices, worker_info.num_workers) + epoch_indices = split_indices[worker_info.id] + return filter( + lambda x: x is not None, # skips the errors + map(self.training_example, epoch_indices), + ) + + def format_error(self, error: Dict[str, Any]) -> str: + return f"On index `{error[self.idx_key]}` got error `{error[self.error_key]}` from source `{error[self.source_key]}`" + + def training_example( + self, + data_idx: DataIndex, + ) -> Optional[Tuple[Data, Data]]: + combined_inputs, combined_outputs = {}, {} + for data_source in self.data_sources: + try: + data_in, data_out = data_source(data_idx) + except EXCEPTIONS as e: + if self.raise_errors: + raise e + self.errors.append( + { + self.error_key: repr(e), + self.source_key: data_source.__name__, + self.idx_key: data_idx, + }, + ) + if self.verbose: + print(self.format_error(self.errors[-1])) + return + for name, modality in data_in.items(): + if name in combined_inputs: + raise ValueError( + f"Multiple DataSources share the same input modality called {name}", + ) + combined_inputs[name] = modality + for name, modality in data_out.items(): + if name in combined_outputs: + raise ValueError( + f"Multiple DataSources share the same output modality called {name}", + ) + combined_outputs[name] = modality + return combined_inputs, combined_outputs diff --git a/ml4ht/data/explore.py b/ml4ht/data/explore.py index 842a32b..ba1701e 100644 --- a/ml4ht/data/explore.py +++ b/ml4ht/data/explore.py @@ -1,12 +1,14 @@ from functools import partial from multiprocessing import Pool, cpu_count -from typing import Callable, List +from typing import Callable, List, TypeVar, Dict import pandas as pd +import numpy as np from ml4ht.data.data_description import DataDescription from ml4ht.data.defines import EXCEPTIONS, SampleID, SampleGetter from ml4ht.data.sample_getter import DataDescriptionSampleGetter +from ml4ht.data.data_source import DataSource, DataIndex ERROR_COL = "error" NO_LOADING_OPTIONS_ERROR = ValueError("No loading options") @@ -15,9 +17,12 @@ LOADING_OPTION_COL = "state" +T = TypeVar("T") + + def build_df( - summarizer: Callable[[SampleID], pd.DataFrame], - sample_ids: List[SampleID], + summarizer: Callable[[T], pd.DataFrame], + sample_ids: List[T], multiprocess_workers: int = 0, ) -> pd.DataFrame: """ @@ -122,13 +127,25 @@ def explore_data_descriptions( ) +def _summarize_tensor(tensor: np.ndarray) -> Dict[str, float]: + out = {} + out["mean"] = tensor.mean() + out["median"] = np.median(tensor) + out["std"] = tensor.std() + out["min"] = tensor.min() + out["max"] = tensor.max() + out["argmax"] = tensor.argmax() + out["shape"] = tensor.shape + return out + + # sample getter explore def _pipeline_sample_getter_summarize_sample_id( sample_id: SampleID, sample_getter: SampleGetter, ) -> pd.DataFrame: """ - Get the summary, dates, states or errors from each TensorMap for a sample id + Get the summary, dates, states or errors from each input and output for a sample id """ out = {SAMPLE_ID_COL: sample_id} try: @@ -137,12 +154,9 @@ def _pipeline_sample_getter_summarize_sample_id( **data[0], **data[1], }.items(): - out[f"{name}_mean"] = [tensor.mean()] - out[f"{name}_std"] = [tensor.std()] - out[f"{name}_min"] = [tensor.min()] - out[f"{name}_max"] = [tensor.max()] - out[f"{name}_argmax"] = [tensor.argmax()] - out[f"{name}_shape"] = [tensor.shape] + tensor_summary = _summarize_tensor(tensor) + for field, val in tensor_summary.items(): + out[f"{name}_{field}"] = [val] except EXCEPTIONS as e: out[ERROR_COL] = [_format_exception(e)] out = pd.DataFrame(out) @@ -155,7 +169,7 @@ def explore_sample_getter( multiprocess_workers: int = 0, ) -> pd.DataFrame: """ - Get the datetime selected, + Summarize the values and errors of a sample getter """ summarize = partial( _pipeline_sample_getter_summarize_sample_id, @@ -166,3 +180,69 @@ def explore_sample_getter( sample_ids, multiprocess_workers, ) + + +def _data_source_auto_summarize_sample_id( + data_idx: DataIndex, + data_source: DataSource, +) -> pd.DataFrame: + out = {f"idx: {k}": [v] for k, v in data_idx.items()} + try: + data_in, data_out = data_source(data_idx) + for name, tensor in data_in.items(): + out.update( + { + f"input: {name}_{field}": [val] + for field, val in _summarize_tensor(tensor).items() + }, + ) + for name, tensor in data_out.items(): + out.update( + { + f"output: {name}_{field}": [val] + for field, val in _summarize_tensor(tensor).items() + }, + ) + except EXCEPTIONS as e: + out[ERROR_COL] = [_format_exception(e)] + return pd.DataFrame(out) + + +def _data_source_summarize_sample_id( + data_idx: DataIndex, + data_source: DataSource, +) -> pd.DataFrame: + out = {f"idx: {k}": [v] for k, v in data_idx.items()} + try: + data_in, data_out = data_source(data_idx) + for name, tensor in data_in.items(): + out[f"input: {name}"] = [tensor] + for name, tensor in data_out.items(): + out[f"output: {name}"] = [tensor] + except EXCEPTIONS as e: + out[ERROR_COL] = [_format_exception(e)] + return pd.DataFrame(out) + + +def explore_data_source( + data_source: DataSource, + data_indices: List[DataIndex], + auto_summarize=True, + multiprocess_workers: int = 0, +) -> pd.DataFrame: + """ + Get DataSource values over all the data indices. + `auto_summarize` converts those values into summary statistics + in the output `DataFrame`. + """ + summarize = partial( + _data_source_auto_summarize_sample_id + if auto_summarize + else _data_source_summarize_sample_id(), + data_source=data_source, + ) + return build_df( + summarize, + data_indices, + multiprocess_workers, + ) diff --git a/ml4ht/data/util/data_source_util.py b/ml4ht/data/util/data_source_util.py new file mode 100644 index 0000000..6638068 --- /dev/null +++ b/ml4ht/data/util/data_source_util.py @@ -0,0 +1,19 @@ +from typing import Sequence + +import pandas as pd + +from ml4ht.data.data_source import DataIndex, range_epoch_idx_generator + + +def df_epoch_index_generator( + df: pd.DataFrame, + shuffle: bool, + training_epochs_per_true_epoch: int, +) -> Sequence[DataIndex]: + epoch_idx_generator = range_epoch_idx_generator( + true_epoch_len=len(df), + shuffle=shuffle, + training_epochs_per_true_epoch=training_epochs_per_true_epoch, + ) + for epoch_indices in epoch_idx_generator: + yield df.iloc[epoch_indices].reindex() diff --git a/tests/data/test_data_source.py b/tests/data/test_data_source.py new file mode 100644 index 0000000..5a592df --- /dev/null +++ b/tests/data/test_data_source.py @@ -0,0 +1,214 @@ +from typing import Iterable +import pytest +import numpy as np +from torch.utils.data import DataLoader +from ml4ht.data.data_source import ( + range_epoch_idx_generator, + data_index_epoch_generator, + TrainingDataset, + TrainingIterableDataset, + DataIndex, +) +from ml4ht.data.data_loader import numpy_collate_fn + + +class TestRangeEpochIdxGenerator: + def test_no_shuffle_true_epoch(self): + epoch_len = 100 + expected_epoch = np.arange(epoch_len) + idx_generator = range_epoch_idx_generator( + true_epoch_len=epoch_len, + shuffle=False, + training_epochs_per_true_epoch=1, + ) + for _ in range(3): + epoch_indices = next(idx_generator) + assert (epoch_indices == expected_epoch).all() + + def test_shuffle_true_epoch(self): + epoch_len = 100 + expected_epoch = np.arange(epoch_len) + idx_generator = range_epoch_idx_generator( + true_epoch_len=epoch_len, + shuffle=True, + training_epochs_per_true_epoch=1, + ) + for _ in range(3): + epoch_indices = next(idx_generator) + sorted_indices = np.sort(epoch_indices) + # did we see the all expected idxs + assert (sorted_indices == expected_epoch).all() + # were they shuffled? + assert (epoch_indices != expected_epoch).any() + + def test_no_shuffle_smaller_epoch(self): + true_epoch_len = 100 + training_steps_per_epoch = 5 + idx_generator = range_epoch_idx_generator( + true_epoch_len=true_epoch_len, + shuffle=False, + training_epochs_per_true_epoch=training_steps_per_epoch, + ) + epoch_len = true_epoch_len // training_steps_per_epoch + for i in range(training_steps_per_epoch * 2): + epoch_indices = next(idx_generator) + start_idx = (i % training_steps_per_epoch) * epoch_len + stop_idx = start_idx + epoch_len + assert (epoch_indices == np.arange(start_idx, stop_idx)).all() + + def test_shuffle_smaller_epoch(self): + true_epoch_len = 100 + training_steps_per_epoch = 6 + idx_generator = range_epoch_idx_generator( + true_epoch_len=true_epoch_len, + shuffle=True, + training_epochs_per_true_epoch=training_steps_per_epoch, + ) + for _ in range(2): + seen_ids = [] + for _ in range(training_steps_per_epoch): + seen_ids += list(next(idx_generator)) + # did we see the all expected idxs + assert sorted(seen_ids) == list(range(true_epoch_len)) + # were they shuffled? + assert seen_ids != list(range(true_epoch_len)) + + +def _input_data_source(idx: DataIndex): + i = np.array(idx["sample_id"]) + if i == 5: + raise ValueError + return {"a": i, "b": -i}, {} + + +def _output_data_source(idx: DataIndex): + i = np.array(idx["sample_id"]) + return {}, {"c": i, "d": -i} + + +def assert_batches_close(expected, actual): + for expected_half, actual_half in zip(expected, actual): + for k, v in expected_half.items(): + np.testing.assert_allclose(actual_half[k], v) + + +class TestTrainingDataset: + @staticmethod + def _get_dset(ids: Iterable[int]): + return TrainingDataset( + [_input_data_source, _output_data_source], + [{"sample_id": i} for i in ids], + ) + + def test_overlapping_inputs(self): + dset = self._get_dset([5]) + with pytest.raises(ValueError): + dset[0] + + @pytest.mark.parametrize( + "multiprocess", + [False, True], + ) + def test_works_in_loader(self, multiprocess): + dset = self._get_dset(range(5)) + loader = DataLoader( + dset, + batch_size=2, + collate_fn=numpy_collate_fn, + num_workers=2 if multiprocess else 0, + drop_last=False, + ) + found_ids = [] + for batch in loader: + sample_ids = batch[0]["a"] + found_ids += list(sample_ids) + assert_batches_close( + numpy_collate_fn( + [ + dset.training_example({"sample_id": sample_id}) + for sample_id in sample_ids + ], + ), + batch, + ) + assert sorted(found_ids) == sorted(range(5)) + + +class TestTrainingIterableDataset: + epoch_generator = data_index_epoch_generator( + [{"sample_id": i} for i in range(6)], + True, + 1, + ) + + def test_skip_errors(self): + dset = TrainingIterableDataset( + [_input_data_source, _output_data_source], + self.epoch_generator, + raise_errors=False, + verbose=True, + ) + inputs = [] + outputs = [] + for x, y in iter(dset): + inputs.append(x) + outputs.append(y) + for i in range(5): + assert _input_data_source({"sample_id": i})[0] in inputs + assert _output_data_source({"sample_id": i})[1] in outputs + assert dset.errors[0] == { + dset.error_key: repr(ValueError()), + dset.idx_key: {"sample_id": 5}, + dset.source_key: _input_data_source.__name__, + } + + def test_raise_errors(self): + dset = TrainingIterableDataset( + [_input_data_source, _output_data_source], + self.epoch_generator, + raise_errors=True, + ) + with pytest.raises(ValueError): + list(iter(dset)) + + def test_overlapping_inputs(self): + dset = TrainingIterableDataset( + [_input_data_source, _input_data_source], + self.epoch_generator, + raise_errors=True, + ) + with pytest.raises(ValueError): + list(iter(dset)) + + @pytest.mark.parametrize( + "multiprocess", + [False, True], + ) + def test_works_in_loader(self, multiprocess): + dset = TrainingIterableDataset( + [_input_data_source, _output_data_source], + self.epoch_generator, + raise_errors=False, + ) + loader = DataLoader( + dset, + batch_size=2, + collate_fn=numpy_collate_fn, + num_workers=2 if multiprocess else 0, + drop_last=False, + ) + found_ids = [] + for batch in loader: + sample_ids = batch[0]["a"] + found_ids += list(sample_ids) + assert_batches_close( + numpy_collate_fn( + [ + dset.training_example({"sample_id": sample_id}) + for sample_id in sample_ids + ], + ), + batch, + ) + assert len(found_ids) == 6 - 1 # right epoch length? + assert set(found_ids) == set(range(5)) # right epoch contents? diff --git a/tests/data/test_explore.py b/tests/data/test_explore.py index 8b44fe1..c743338 100644 --- a/tests/data/test_explore.py +++ b/tests/data/test_explore.py @@ -19,8 +19,12 @@ _pipeline_sample_getter_summarize_sample_id, build_df, _format_exception, + _summarize_tensor, + _data_source_summarize_sample_id, + _data_source_auto_summarize_sample_id, ) from ml4ht.data.sample_getter import DataDescriptionSampleGetter +from ml4ht.data.data_source import DataIndex RAW_DATA_1 = { 0: { @@ -161,3 +165,46 @@ def test_fail_one_data_description(self): def test_no_loading_options(self): row = _pipeline_sample_getter_summarize_sample_id(4, PIPE).iloc[0] assert row[ERROR_COL] == _format_exception(NoDatesAvailableError()) + + +def _data_source(idx: DataIndex): + i = np.array(idx["sample_id"]) + if i == 5: + raise ValueError + return {"a": i, "b": -i}, {"c": np.array(i ** 2)} + + +class TestExploreDataSource: + def test_data_source_summarize_sample_id(self): + idx = {"sample_id": 1} + out = _data_source_summarize_sample_id(idx, _data_source) + assert out["idx: sample_id"].iloc[0] == 1 + expected_in, expected_out = _data_source(idx) + for k, v in expected_in.items(): + assert out[f"input: {k}"].iloc[0] == v + for k, v in expected_out.items(): + assert out[f"output: {k}"].iloc[0] == v + + def test_data_source_summarize_sample_id_error(self): + idx = {"sample_id": 5} + out = _data_source_summarize_sample_id(idx, _data_source) + assert out["idx: sample_id"].iloc[0] == 5 + assert out[ERROR_COL].iloc[0] == _format_exception(ValueError()) + + def test_data_source_auto_summarize_sample_id(self): + idx = {"sample_id": 1} + out = _data_source_auto_summarize_sample_id(idx, _data_source) + assert out["idx: sample_id"].iloc[0] == 1 + expected_in, expected_out = _data_source(idx) + for k, v in expected_in.items(): + for field, val in _summarize_tensor(v).items(): + assert out[f"input: {k}_{field}"].iloc[0] == val + for k, v in expected_out.items(): + for field, val in _summarize_tensor(v).items(): + assert out[f"output: {k}_{field}"].iloc[0] == val + + def test_data_source_auto_summarize_sample_id_error(self): + idx = {"sample_id": 5} + out = _data_source_summarize_sample_id(idx, _data_source) + assert out["idx: sample_id"].iloc[0] == 5 + assert out[ERROR_COL].iloc[0] == _format_exception(ValueError()) diff --git a/tests/data/util/test_data_source_util.py b/tests/data/util/test_data_source_util.py new file mode 100644 index 0000000..14457d9 --- /dev/null +++ b/tests/data/util/test_data_source_util.py @@ -0,0 +1,29 @@ +import pandas as pd + +from ml4ht.data.util.data_source_util import df_epoch_index_generator + + +class TestDfEpochIndexGenerator: + df = pd.DataFrame( + { + "a": [1, 2, 3], + "b": ["1", "2", "3"], + }, + ) + + def test_shuffle_smaller_epoch(self): + training_steps_per_epoch = 2 + idx_generator = df_epoch_index_generator( + self.df, + shuffle=True, + training_epochs_per_true_epoch=training_steps_per_epoch, + ) + for _ in range(2): + seen_idxs = [] + for _ in range(training_steps_per_epoch): + seen_idxs.append(next(idx_generator)) + print(pd.concat(seen_idxs)) + pd.testing.assert_frame_equal( + pd.concat(seen_idxs).sort_values("a"), + self.df, + )