Skip to content

ENH: Add options to consolidate with ecg2x notebooks #28

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

ENH: Add options to consolidate with ecg2x notebooks #28

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
aec9e93
ENH: Add options to consolidate with ecg2x notebooks
daniellepace May 1, 2024
90d369e
STYLE: Pass style checks
daniellepace Aug 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 14 additions & 3 deletions ml4ht/data/data_loader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -282,21 +282,32 @@ def __getitem__(self, item: int) -> Batch:
)


def numpy_collate_fn(samples: List[Batch]) -> Batch:
def numpy_collate_fn(
samples: List[Batch],
auto_float: bool = True,
) -> Batch:
"""
Merges a list of ml4ht batch formatted data.
Can be used as 'collate_fn` in torch.utils.data.DataLoader
so that the torch data loader is compatible with tensorflow models
"""
# construct correctly-shaped empty arrays for input and output of model
in_batch_keys = list(samples[0][0])
if auto_float:
in_dtypes = {k: np.float32 for k in in_batch_keys}
else:
in_dtypes = {k: samples[0][0][k].dtype for k in in_batch_keys}
in_batch = {
k: np.empty((len(samples),) + samples[0][0][k].shape, dtype=np.float32)
k: np.empty((len(samples),) + samples[0][0][k].shape, dtype=in_dtypes[k])
for k in in_batch_keys
}
out_batch_keys = list(samples[0][1])
if auto_float:
out_dtypes = {k: np.float32 for k in out_batch_keys}
else:
out_dtypes = {k: samples[0][1][k].dtype for k in out_batch_keys}
out_batch = {
k: np.empty((len(samples),) + samples[0][1][k].shape, dtype=np.float32)
k: np.empty((len(samples),) + samples[0][1][k].shape, dtype=out_dtypes[k])
for k in out_batch_keys
}
# fill in the values of the input and output arrays
Expand Down
4 changes: 4 additions & 0 deletions ml4ht/data/sample_getter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,10 +29,12 @@ def __init__(
input_data_descriptions: List[DataDescription],
output_data_descriptions: List[DataDescription],
option_picker: OptionPicker = None,
restricted_sample_id_idx=None,
):
self.input_data_descriptions = input_data_descriptions
self.output_data_descriptions = output_data_descriptions
self.option_picker = option_picker or self._default_option_picker
self.restricted_sample_id_idx = restricted_sample_id_idx

@staticmethod
def _default_option_picker(
Expand Down Expand Up @@ -67,6 +69,8 @@ def __call__(self, sample_id: SampleID) -> Batch:
sample_id,
self.input_data_descriptions + self.output_data_descriptions,
)
if self.restricted_sample_id_idx is not None:
sample_id = sample_id[self.restricted_sample_id_idx]
tensors_in = self._half_batch(sample_id, loading_options, True)
tensors_out = self._half_batch(sample_id, loading_options, False)
return tensors_in, tensors_out
Loading