def testChooseFastestErrorWithIncompatibleInput(self, slices_a, slices_b,
error_msg):
dataset_a = dataset_ops.Dataset.from_tensor_slices(slices_a)
dataset_b = dataset_ops.Dataset.from_tensor_slices(slices_b)
# The error is raised at dataset creation time.
if context.executing_eagerly():
with self.assertRaises(errors.InvalidArgumentError):
merge = optimization._ChooseFastestDataset([dataset_a, dataset_b])
else:
merge = optimization._ChooseFastestDataset([dataset_a, dataset_b])
self.assertDatasetProduces(
merge, expected_error=(errors.InvalidArgumentError, error_msg))
def testChooseFastestErrorWithIncompatibleInput(self, slices_a, slices_b,
error_msg):
dataset_a = dataset_ops.Dataset.from_tensor_slices(slices_a)
dataset_b = dataset_ops.Dataset.from_tensor_slices(slices_b)
# The error is raised at dataset creation time.
if context.executing_eagerly():
with self.assertRaises(errors.InvalidArgumentError):
merge = optimization._ChooseFastestDataset([dataset_a, dataset_b])
else:
merge = optimization._ChooseFastestDataset([dataset_a, dataset_b])
self.assertDatasetProduces(
merge, expected_error=(errors.InvalidArgumentError, error_msg))
def testChooseFastestManyInputs(self):
dataset = dataset_ops.Dataset.from_tensor_slices([0, 1, 2, 3, 4])
merge = optimization._ChooseFastestDataset([dataset for _ in range(5)])
self.assertDatasetProduces(
merge,
expected_output=[0, 1, 2, 3, 4],
expected_shapes=dataset.output_shapes)
def testChooseFastestSimple(self):
dataset = dataset_ops.Dataset.from_tensor_slices([0, 1, 2, 3, 4])
merge = optimization._ChooseFastestDataset([dataset, dataset])
self.assertDatasetProduces(
merge,
expected_output=[0, 1, 2, 3, 4],
expected_shapes=dataset.output_shapes)
def testChooseFastestManyInputs(self):
dataset = dataset_ops.Dataset.from_tensor_slices([0, 1, 2, 3, 4])
merge = optimization._ChooseFastestDataset([dataset for _ in range(5)])
self.assertDatasetProduces(
merge,
expected_output=[0, 1, 2, 3, 4],
expected_shapes=dataset_ops.get_legacy_output_shapes(dataset))
def build_ds():
dataset = dataset_ops.Dataset.range(num_outputs)
map_fn = lambda x: x * 2
return optimization._ChooseFastestDataset([ # pylint: disable=protected-access
dataset.map(map_fn).batch(batch_size),
dataset.batch(batch_size).map(map_fn)
])
def build_ds():
dataset = dataset_ops.Dataset.range(num_outputs)
map_fn = lambda x: x * 2
return optimization._ChooseFastestDataset([ # pylint: disable=protected-access
dataset.map(map_fn).batch(batch_size),
dataset.batch(batch_size).map(map_fn)
])
def testChooseFastest(self):
dataset = dataset_ops.Dataset.range(600)
f = lambda x: 2 * x
dataset_a = dataset.batch(50).map(f)
dataset_b = dataset.map(f).batch(50)
merge = optimization._ChooseFastestDataset([dataset_a, dataset_b])
self.assertDatasetProduces(
merge,
expected_output=[[i * 2 for i in range(j * 50, (j + 1) * 50)]
for j in range(12)],
expected_shapes=dataset_ops.get_legacy_output_shapes(dataset_a))
def testChooseFastest(self):
dataset = dataset_ops.Dataset.range(600)
f = lambda x: 2 * x
dataset_a = dataset.batch(50).map(f)
dataset_b = dataset.map(f).batch(50)
merge = optimization._ChooseFastestDataset([dataset_a, dataset_b])
self.assertDatasetProduces(
merge,
expected_output=[
[i * 2 for i in range(j * 50, (j + 1) * 50)] for j in range(12)
],
expected_shapes=dataset_ops.get_legacy_output_shapes(dataset_a))
def benchmark_choose_fastest(self):
dataset = dataset_ops.Dataset.range(1000**2).repeat()
options = dataset_ops.Options()
options.experimental_optimization.apply_default_optimizations = False
dataset = dataset.with_options(options)
map_batch_dataset = dataset.map(lambda x: x + 1).batch(100)
batch_map_dataset = dataset.batch(100).map(lambda x: x + 1)
merge_dataset = optimization._ChooseFastestDataset( # pylint: disable=protected-access
[batch_map_dataset, map_batch_dataset])
self._benchmark(dataset=map_batch_dataset, name="map_batch_dataset")
self._benchmark(dataset=batch_map_dataset, name="batch_map_dataset")
self._benchmark(dataset=merge_dataset, name="merge_dataset")
def benchmarkChooseFastest(self):
dataset = dataset_ops.Dataset.range(1000**2).repeat()
options = dataset_ops.Options()
options.experimental_optimization.apply_default_optimizations = False
dataset = dataset.with_options(options)
map_batch_dataset = dataset.map(lambda x: x + 1).batch(100)
batch_map_dataset = dataset.batch(100).map(lambda x: x + 1)
merge_dataset = optimization._ChooseFastestDataset( # pylint: disable=protected-access
[batch_map_dataset, map_batch_dataset])
self._benchmark(map_batch_dataset, "map_batch_dataset")
self._benchmark(batch_map_dataset, "batch_map_dataset")
self._benchmark(merge_dataset, "merge_dataset")
