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")