def _init(self, num_processes):
self._x = [ivy.array([0, 1]), ivy.array([2, 3, 4, 5, 6, 7, 8, 9])]
dataset_container = ivy.Container({'x': self._x})
dataset = Dataset(dataset_container,
'base',
dataset_container.shape[0],
num_processes=num_processes)
dataset = dataset.unbatch('unbatched', num_processes=num_processes)
self._dataset = dataset.batch('batched',
3,
num_processes=num_processes)
def _init(self, array_shape, num_processes):
x = [
ivy.array([[0], [1], [2]]),
ivy.array([[3], [4], [5]]),
ivy.array([[6], [7], [8]])
]
self._x = [ivy.reshape(item, array_shape) for item in x]
dataset_container = ivy.Container({'x': x})
dataset = Dataset(dataset_container,
'base',
dataset_container.shape[0],
num_processes=num_processes)
self._dataset = dataset.unbatch('unbatched',
num_processes=num_processes)
def _init(self, array_shape, num_processes):
x = [
ivy.array(0),
ivy.array(1),
ivy.array(2),
ivy.array(3),
ivy.array(4),
ivy.array(5),
ivy.array(6),
ivy.array(7),
ivy.array(8)
]
self._x = [ivy.reshape(item, array_shape) for item in x]
dataset_container = ivy.Container({'x': self._x})
self._dataset = Dataset(dataset_container,
'base',
dataset_container.shape[0],
num_processes=num_processes)
def _init(self, array_shape, num_processes):
x = [
ivy.array(0),
ivy.array(1),
ivy.array(2),
ivy.array(3),
ivy.array(4),
ivy.array(5),
ivy.array(6),
ivy.array(7),
ivy.array(8),
ivy.array(9)
]
self._x = [ivy.reshape(item, array_shape) for item in x]
def sleep_fn(cont):
start_time = time.perf_counter()
while True:
if time.perf_counter() - start_time > 0.011:
return cont
dataset_container = ivy.Container({'x': self._x})
# without pre-fetch
dataset_wo_prefetch = Dataset(copy.deepcopy(dataset_container),
'base',
dataset_container.shape[0],
with_caching=False,
cache_size=0,
num_processes=num_processes)
self._dataset_wo_prefetch = dataset_wo_prefetch.map('sleep', sleep_fn)
# with pre-fetch
dataset_w_prefetch = Dataset(copy.deepcopy(dataset_container),
'base',
dataset_container.shape[0],
with_caching=False,
cache_size=0,
num_processes=num_processes)
dataset_w_prefetch = dataset_w_prefetch.map('sleep', sleep_fn)
self._dataset_w_prefetch = dataset_w_prefetch.prefetch('prefetch', 1)
def _get_dataset(self, starting_example, ending_example):
class ContainerIdxMap:
def __init__(self,
sizes,
fpath_template=None,
seq_idxs=None,
start=None,
end=None,
max_seq_len=None,
conts_to_skip=None,
pruned_sizes=None):
if isinstance(sizes, (tuple, list)):
pruned_sizes = ivy.default(pruned_sizes, [
SeqDataLoader._compute_seq_len(i, sl, conts_to_skip)
for i, sl in enumerate(sizes)
])
num_empty = sum([ps == 0 for ps in pruned_sizes])
self._raw_sizes = dict(
zip(range(start, end + 1 + num_empty),
sizes[start:end + 1 + num_empty]))
self._pruned_sizes = dict(
zip(range(start, end + 1 + num_empty),
pruned_sizes[start:end + 1 + num_empty]))
elif isinstance(sizes, (int, dict)):
self._raw_sizes = sizes
self._pruned_sizes = ivy.default(pruned_sizes, sizes)
if isinstance(self._pruned_sizes, int):
pruned_dict = dict()
for seq_idx, win_idx in conts_to_skip:
if seq_idx not in pruned_dict:
pruned_dict[seq_idx] = list()
pruned_dict[seq_idx].append(win_idx)
pruned_dict = {
k: len(set(v))
for k, v in pruned_dict.items()
}
pruned_sizes_dict = {
k: self._pruned_sizes - num_pruned
for k, num_pruned in pruned_dict.items()
}
num_empty = sum(
[size == 0 for size in pruned_sizes_dict.values()])
pruned_sizes = collections.defaultdict(
lambda: self._pruned_sizes, pruned_sizes_dict)
else:
num_empty = sum([ps == 0 for ps in self._pruned_sizes])
else:
raise Exception(
'Invalid type for sizes, expected one of int, dict, tuple or list,'
'but found {} or type {}'.format(sizes, type(sizes)))
self._constant_size = isinstance(self._raw_sizes, int)
if max_seq_len:
self._max_seq_len = max_seq_len
else:
self._max_seq_len = self._pruned_sizes if self._constant_size else max(
self._pruned_sizes.values())
self._fpath_template = fpath_template
self._conts_to_skip = conts_to_skip
if seq_idxs:
self._seq_idxs = seq_idxs
else:
vals = [
v
for i, v in enumerate(range(start, end + 1 +
num_empty))
if pruned_sizes[i] > 0
]
keys = range(0, min(end - start + 1 + num_empty,
len(vals)))
self._seq_idxs = dict(zip(keys, vals))
def __getitem__(self, slice_obj):
if isinstance(slice_obj, slice):
seq_idxs = collections.OrderedDict([
(i, self._seq_idxs[idx]) for i, idx in enumerate(
range(slice_obj.start, slice_obj.stop,
ivy.default(slice_obj.step, 1)))
])
elif isinstance(slice_obj, int):
seq_idxs = collections.OrderedDict(
{0: self._seq_idxs[slice_obj]})
else:
raise Exception(
'Invalid type for slice_obj, expected either slice or int,'
'but found {} of type {}'.format(
slice_obj, type(slice_obj)))
if self._constant_size:
sizes = self._raw_sizes
else:
sizes = collections.OrderedDict({
seq_idx: self._raw_sizes[seq_idx]
for seq_idx in seq_idxs.values()
})
return ContainerIdxMap(sizes,
self._fpath_template,
seq_idxs,
max_seq_len=self._max_seq_len,
conts_to_skip=self._conts_to_skip,
pruned_sizes=self._pruned_sizes)
def __len__(self):
return len(self._seq_idxs)
def shuffle(self):
mapped_idxs = list(self._seq_idxs.values())
np.random.shuffle(mapped_idxs)
self._seq_idxs = collections.OrderedDict(
zip(self._seq_idxs.keys(), mapped_idxs))
def to_idxs(self):
seq_idxs = self._seq_idxs.values()
sizes = [
self._raw_sizes
if self._constant_size else self._raw_sizes[seq_idx]
for seq_idx in seq_idxs
]
rets = [[(seq_idx, win_idx) for win_idx in range(size)
if not SeqDataLoader._skip_cont(
seq_idx, win_idx, self._conts_to_skip)]
for seq_idx, size in zip(seq_idxs, sizes)]
return [
r + [(None, None)] * (self._max_seq_len - len(r))
for r in rets if list(set(r)) != [None]
]
def to_filepaths(self):
if not ivy.exists(self._fpath_template):
raise Exception(
'to_filepaths method is not valid if fpath_template has not been specified'
'in the constructor.')
seq_idxs = self._seq_idxs.values()
sizes = [
self._raw_sizes
if self._constant_size else self._raw_sizes[seq_idx]
for seq_idx in seq_idxs
]
rets = [[
self._fpath_template % (seq_idx, win_idx)
for win_idx in range(size) if not SeqDataLoader._skip_cont(
seq_idx, win_idx, self._conts_to_skip)
] for seq_idx, size in zip(seq_idxs, sizes)]
return [
r + [''] * (self._max_seq_len - len(r)) for r in rets
if ''.join(r) != ''
]
@property
def sizes(self):
return self._pruned_sizes
# container filepaths
if self._spec.container_load_mode in ['preload', 'dynamic']:
fpath_template = os.path.join(
self._container_data_dir,
self._spec.dataset_spec.cont_fname_template)
else:
fpath_template = None
container_idx_map = ContainerIdxMap(
self._spec.dataset_spec.unpruned_sequence_lengths,
fpath_template,
start=starting_example,
end=ending_example,
conts_to_skip=self._spec.containers_to_skip)
if self._spec.num_sequences != -1:
container_idx_map = container_idx_map[0:self._spec.num_sequences]
# shuffle sequences
if self._spec.preshuffle_data:
container_idx_map.shuffle()
# extract sequence lengths
if self._fixed_sequence_length:
self._sequence_lengths =\
collections.OrderedDict(zip(range(len(container_idx_map)),
[self._spec.dataset_spec.sequence_lengths] * len(container_idx_map)))
self._windows_per_seq = self._sequence_lengths[
0] - self._window_size + 1
# windowing values
window_idxs_per_seq = ivy.reshape(
ivy.arange(self._windows_per_seq, 0, 1),
(self._windows_per_seq, 1))
gather_idxs_list = list()
for x in window_idxs_per_seq:
gather_idxs_list.append(
ivy.expand_dims(
ivy.arange(x[0] + self._window_size, x[0], 1), 0))
gather_idxs = ivy.concatenate(gather_idxs_list, 0)
self._gather_idxs = \
ivy.to_numpy(ivy.reshape(gather_idxs, (self._windows_per_seq * self._window_size, 1))).tolist()
else:
self._sequence_lengths = container_idx_map.sizes
# maybe pre-load containers
if self._spec.container_load_mode == 'preload':
# load containers with vector data and image filepath entries
container_slices = self._get_containers_w_filepath_img_entries_as_tensor_slices(
container_idx_map.to_filepaths())
if self._first_frame_validity_fn is not None:
container_slices =\
self._first_frame_validity_fn(container_slices, [ending_example - starting_example + 1])
# prune unwanted chains of keys
if 'unused_key_chains' in self._spec:
container_slices = self._prune_unused_key_chains(
container_slices)
dataset = Dataset(ivy.Container.list_stack([
c[0]
for c in container_slices.unstack(0, container_slices.shape[0])
], 0),
'base',
container_slices.shape[0],
numpy_loading=True,
cache_size=self._base_cache_size,
queue_timeout=self._spec.queue_timeout)
else:
if self._spec.container_load_mode == 'dynamic':
# load containers with filepath entries
dataset = Dataset(ivy.Container({'fpaths': container_idx_map}),
'base',
len(container_idx_map),
trans_fn=lambda cont: cont.map(
lambda x_, kc: x_.to_filepaths()),
elementwise_query_fn=False,
numpy_loading=True,
cache_size=self._base_cache_size,
queue_timeout=self._spec.queue_timeout)
dataset = dataset.map('loaded_json', self._load_json_files,
self._num_workers.loaded_json)
dataset = dataset.map('parsed_json', self._parse_json_strings,
self._num_workers.parsed_json)
else:
dataset = Dataset(ivy.Container({'idx_map':
container_idx_map}),
'base',
len(container_idx_map),
trans_fn=lambda cont: self._spec.
custom_container_load_fn(self, cont),
elementwise_query_fn=False,
numpy_loading=True,
cache_size=self._base_cache_size,
queue_timeout=self._spec.queue_timeout)
if 'unused_key_chains' in self._spec:
dataset = dataset.map('keychain_pruned',
self._prune_unused_key_chains,
self._num_workers.keychain_pruned)
if self._first_frame_validity_fn is not None:
dataset = dataset.map(
'valid_first_frames',
lambda x_: self._first_frame_validity_fn(x_, None),
self._num_workers.valid_first_frames)
if not (self._spec.dataset_spec.sequence_lengths == 1
and self._window_size == 1):
# ToDo: add other conditionals which make the loading more efficient if only one of the
# above two conditions is True
dataset = dataset.map(
'windowed', self._group_container_into_windowed_container,
self._num_workers.windowed)
dataset = dataset.unbatch(
'unbatched',
self._num_workers.unbatched,
batch_sizes=[
max(seq_len, self._window_size) - self._window_size + 1
for seq_len in self._sequence_lengths.values()
if seq_len > 0
])
if self._spec.shuffle_buffer_size > 0:
dataset = dataset.shuffle('shuffled',
self._spec.shuffle_buffer_size,
self._num_workers.shuffled)
dataset = dataset.map('loaded_data',
self._load_data_from_filepath_tensors,
self._num_workers.loaded_data)
dataset = dataset.batch('batched', self._batch_size,
self._num_workers.batched)
dataset = dataset.map(
'from_np',
lambda cont: cont.map(lambda x_, kc: ivy.array(x_, dev_str='cpu')),
self._num_workers.from_np,
numpy_loading=False)
if ivy.exists(self._spec.post_proc_fn):
dataset = dataset.map('post_processed', self._spec.post_proc_fn,
self._num_workers.post_processed)
if self._spec.with_prefetching:
dataset = dataset.prefetch('prefetch')
# ToDo: find way to make pre-fetching to GPU actually pre-fetch, ideally using multi-processing.
# For example, swapping prefetch and to_gpu ops around would work if to_gpu could accept self._num_workers.
if self._spec.prefetch_to_devs:
if isinstance(self._spec.prefetch_to_devs, str):
dataset = dataset.to_dev('to_dev', self._spec.prefetch_to_devs)
elif len(self._spec.prefetch_to_devs) == 1:
dataset = dataset.to_dev('to_dev',
self._spec.prefetch_to_devs[0])
else:
dataset = dataset.to_devs('to_devs',
self._spec.prefetch_to_devs)
return dataset
class TestShuffle:
def _init(self, array_shape, num_processes):
x = [
ivy.array(0),
ivy.array(1),
ivy.array(2),
ivy.array(3),
ivy.array(4),
ivy.array(5),
ivy.array(6),
ivy.array(7),
ivy.array(8)
]
self._x = [ivy.reshape(item, array_shape) for item in x]
dataset_container = ivy.Container({'x': self._x})
self._dataset = Dataset(dataset_container,
'base',
dataset_container.shape[0],
num_processes=num_processes).shuffle(
'shuffled', 9)
@pytest.mark.parametrize("array_shape", [[1], []])
@pytest.mark.parametrize("num_processes", [1, 2])
def test_single(self, dev_str, f, call, array_shape, num_processes):
if call is helpers.mx_call and num_processes == 2:
pytest.skip()
ivy.seed(0)
np.random.seed(0)
self._init(array_shape, num_processes)
assert list(self._dataset[0].x.shape) == array_shape
assert list(self._dataset[4].x.shape) == array_shape
assert list(self._dataset[8].x.shape) == array_shape
check0 = not np.allclose(ivy.to_numpy(self._dataset[0].x),
ivy.to_numpy(self._x[0]))
check1 = not np.allclose(ivy.to_numpy(self._dataset[1].x),
ivy.to_numpy(self._x[1]))
check2 = not np.allclose(ivy.to_numpy(self._dataset[2].x),
ivy.to_numpy(self._x[2]))
check3 = not np.allclose(ivy.to_numpy(self._dataset[3].x),
ivy.to_numpy(self._x[3]))
check4 = not np.allclose(ivy.to_numpy(self._dataset[4].x),
ivy.to_numpy(self._x[4]))
check5 = not np.allclose(ivy.to_numpy(self._dataset[5].x),
ivy.to_numpy(self._x[5]))
check6 = not np.allclose(ivy.to_numpy(self._dataset[6].x),
ivy.to_numpy(self._x[6]))
check7 = not np.allclose(ivy.to_numpy(self._dataset[7].x),
ivy.to_numpy(self._x[7]))
check8 = not np.allclose(ivy.to_numpy(self._dataset[8].x),
ivy.to_numpy(self._x[8]))
assert check0 or check1 or check2 or check3 or check4 or check5 or check6 or check7 or check8
# close
self._dataset.close()
del self._dataset
@pytest.mark.parametrize("array_shape", [[1], []])
@pytest.mark.parametrize("num_processes", [1, 2])
def test_single_wrapped(self, dev_str, f, call, array_shape,
num_processes):
if call is helpers.mx_call and num_processes == 2:
pytest.skip()
ivy.seed(0)
np.random.seed(0)
self._init(array_shape, num_processes)
assert list(self._dataset[9].x.shape) == array_shape
assert list(self._dataset[11].x.shape) == array_shape
assert list(self._dataset[-1].x.shape) == array_shape
assert list(self._dataset[-2].x.shape) == array_shape
check0 = not np.allclose(ivy.to_numpy(self._dataset[9].x),
ivy.to_numpy(self._x[0]))
check1 = not np.allclose(ivy.to_numpy(self._dataset[11].x),
ivy.to_numpy(self._x[2]))
check2 = not np.allclose(ivy.to_numpy(self._dataset[-1].x),
ivy.to_numpy(self._x[-1]))
check3 = not np.allclose(ivy.to_numpy(self._dataset[-2].x),
ivy.to_numpy(self._x[-2]))
assert check0 or check1 or check2 or check3
# close
self._dataset.close()
del self._dataset
@pytest.mark.parametrize("array_shape", [[1], []])
@pytest.mark.parametrize("num_processes", [1, 2])
def test_slice(self, dev_str, f, call, array_shape, num_processes):
if call is helpers.mx_call and num_processes == 2:
pytest.skip()
ivy.seed(0)
np.random.seed(0)
self._init(array_shape, num_processes)
assert len(self._dataset[0:3].x) == 3
assert list(self._dataset[0:3].x[0].shape) == array_shape
assert len(self._dataset[3:6].x) == 3
assert list(self._dataset[3:6].x[0].shape) == array_shape
assert len(self._dataset[6:9].x) == 3
assert list(self._dataset[6:9].x[0].shape) == array_shape
check0 = not np.allclose(ivy.to_numpy(self._dataset[0:3].x[0]),
ivy.to_numpy(self._x[0]))
check1 = not np.allclose(ivy.to_numpy(self._dataset[3:6].x[1]),
ivy.to_numpy(self._x[4]))
check2 = not np.allclose(ivy.to_numpy(self._dataset[6:9].x[0]),
ivy.to_numpy(self._x[6]))
check3 = not np.allclose(ivy.to_numpy(self._dataset[6:9].x[1]),
ivy.to_numpy(self._x[7]))
assert check0 or check1 or check2 or check3
# close
self._dataset.close()
del self._dataset
@pytest.mark.parametrize("array_shape", [[1], []])
@pytest.mark.parametrize("num_processes", [1, 2])
def test_slice_wrapped(self, dev_str, f, call, array_shape, num_processes):
if call is helpers.mx_call and num_processes == 2:
pytest.skip()
ivy.seed(0)
np.random.seed(0)
self._init(array_shape, num_processes)
assert len(self._dataset[-1:1].x) == 2
assert list(self._dataset[-1:1].x[0].shape) == array_shape
assert len(self._dataset[-1:1].x) == 2
assert list(self._dataset[-1:1].x[1].shape) == array_shape
assert len(self._dataset[9:11].x) == 2
assert list(self._dataset[9:11].x[0].shape) == array_shape
assert len(self._dataset[9:11].x) == 2
assert list(self._dataset[9:11].x[0].shape) == array_shape
check0 = not np.allclose(ivy.to_numpy(self._dataset[-1:1].x[0]),
ivy.to_numpy(self._x[8]))
check1 = not np.allclose(ivy.to_numpy(self._dataset[-1:1].x[1]),
ivy.to_numpy(self._x[0]))
check2 = not np.allclose(ivy.to_numpy(self._dataset[9:11].x[0]),
ivy.to_numpy(self._x[0]))
check3 = not np.allclose(ivy.to_numpy(self._dataset[9:11].x[1]),
ivy.to_numpy(self._x[1]))
assert check0 or check1 or check2 or check3
# close
self._dataset.close()
del self._dataset
class TestQueries:
def _init(self, array_shape, num_processes):
x = [
ivy.array(0),
ivy.array(1),
ivy.array(2),
ivy.array(3),
ivy.array(4),
ivy.array(5),
ivy.array(6),
ivy.array(7),
ivy.array(8)
]
self._x = [ivy.reshape(item, array_shape) for item in x]
dataset_container = ivy.Container({'x': self._x})
self._dataset = Dataset(dataset_container,
'base',
dataset_container.shape[0],
num_processes=num_processes)
@pytest.mark.parametrize("array_shape", [[1], []])
@pytest.mark.parametrize("num_processes", [1, 2])
def test_single(self, dev_str, f, call, array_shape, num_processes):
if call is helpers.mx_call and num_processes == 2:
pytest.skip()
self._init(array_shape, num_processes)
assert list(self._dataset[0].x.shape) == array_shape
assert list(self._dataset[4].x.shape) == array_shape
assert list(self._dataset[8].x.shape) == array_shape
assert np.allclose(ivy.to_numpy(self._dataset[0].x),
ivy.to_numpy(self._x[0]))
assert np.allclose(ivy.to_numpy(self._dataset[4].x),
ivy.to_numpy(self._x[4]))
assert np.allclose(ivy.to_numpy(self._dataset[8].x),
ivy.to_numpy(self._x[8]))
# close
self._dataset.close()
del self._dataset
@pytest.mark.parametrize("array_shape", [[1], []])
@pytest.mark.parametrize("num_processes", [1, 2])
def test_single_wrapped(self, dev_str, f, call, array_shape,
num_processes):
if call is helpers.mx_call and num_processes == 2:
pytest.skip()
self._init(array_shape, num_processes)
assert list(self._dataset[9].x.shape) == array_shape
assert list(self._dataset[11].x.shape) == array_shape
assert list(self._dataset[-1].x.shape) == array_shape
assert list(self._dataset[-2].x.shape) == array_shape
assert np.allclose(ivy.to_numpy(self._dataset[9].x),
ivy.to_numpy(self._x[0]))
assert np.allclose(ivy.to_numpy(self._dataset[11].x),
ivy.to_numpy(self._x[2]))
assert np.allclose(ivy.to_numpy(self._dataset[-1].x),
ivy.to_numpy(self._x[-1]))
assert np.allclose(ivy.to_numpy(self._dataset[-2].x),
ivy.to_numpy(self._x[-2]))
# close
self._dataset.close()
del self._dataset
@pytest.mark.parametrize("array_shape", [[1], []])
@pytest.mark.parametrize("num_processes", [1, 2])
def test_slice(self, dev_str, f, call, array_shape, num_processes):
if call is helpers.mx_call and num_processes == 2:
pytest.skip()
self._init(array_shape, num_processes)
assert len(self._dataset[0:3].x) == 3
assert list(self._dataset[0:3].x[0].shape) == array_shape
assert len(self._dataset[3:6].x) == 3
assert list(self._dataset[3:6].x[0].shape) == array_shape
assert len(self._dataset[6:9].x) == 3
assert list(self._dataset[6:9].x[0].shape) == array_shape
assert np.allclose(ivy.to_numpy(self._dataset[0:3].x[0]),
ivy.to_numpy(self._x[0]))
assert np.allclose(ivy.to_numpy(self._dataset[3:6].x[1]),
ivy.to_numpy(self._x[4]))
assert np.allclose(ivy.to_numpy(self._dataset[6:9].x[2]),
ivy.to_numpy(self._x[8]))
# close
self._dataset.close()
del self._dataset
@pytest.mark.parametrize("array_shape", [[1], []])
@pytest.mark.parametrize("num_processes", [1, 2])
def test_slice_wrapped(self, dev_str, f, call, array_shape, num_processes):
if call is helpers.mx_call and num_processes == 2:
pytest.skip()
self._init(array_shape, num_processes)
assert len(self._dataset[-1:1].x) == 2
assert list(self._dataset[-1:1].x[0].shape) == array_shape
assert len(self._dataset[-1:1].x) == 2
assert list(self._dataset[-1:1].x[1].shape) == array_shape
assert len(self._dataset[9:11].x) == 2
assert list(self._dataset[9:11].x[0].shape) == array_shape
assert len(self._dataset[9:11].x) == 2
assert list(self._dataset[9:11].x[0].shape) == array_shape
assert np.allclose(ivy.to_numpy(self._dataset[-1:1].x[0]),
ivy.to_numpy(self._x[8]))
assert np.allclose(ivy.to_numpy(self._dataset[-1:1].x[1]),
ivy.to_numpy(self._x[0]))
assert np.allclose(ivy.to_numpy(self._dataset[9:11].x[0]),
ivy.to_numpy(self._x[0]))
assert np.allclose(ivy.to_numpy(self._dataset[9:11].x[1]),
ivy.to_numpy(self._x[1]))
# close
self._dataset.close()
del self._dataset