def _apply_fn(dataset): options = dataset_ops.Options() options.experimental_autotune = False opt_options = optimization_options.OptimizationOptions() opt_options.apply_default_optimizations = False options.experimental_optimization = opt_options return _CopyToDeviceDataset( dataset, target_device=target_device, source_device=source_device).with_options(options)
def testOptionsHaveDefaults(self): options1 = dataset_ops.Options() options2 = dataset_ops.Options() self.assertIsNot(options1.experimental_optimization, options2.experimental_optimization) self.assertIsNot(options1.threading, options2.threading) self.assertEqual(options1.experimental_optimization, optimization_options.OptimizationOptions()) self.assertEqual(options1.threading, threading_options.ThreadingOptions())
def _make_dataset(node_names): dataset = base_dataset.apply(optimization.assert_next(node_names)) dataset = dataset.map(map_fn, num_parallel_calls) dataset = dataset.batch(100) options = dataset_ops.Options() opt_options = optimization_options.OptimizationOptions() opt_options.map_and_batch_fusion = False options.experimental_optimization = opt_options dataset = dataset.with_options(options) return dataset
def testOptimizationDisableDefault(self): """Tests that we can disable all static optimizations enabled by default. If the `apply_default_optimizations` optimization options flag is False, only explicitly enabled optimizations will be applied. """ options = dataset_ops.Options() opt_options = optimization_options.OptimizationOptions() opt_options.hoist_random_uniform = True opt_options.apply_default_optimizations = False options.experimental_optimization = opt_options expected_optimizations = ["hoist_random_uniform"] self.assertEqual(options._static_optimizations(), expected_optimizations)
def _compare(self, input_dataset, map_fn, batch_size, input_size, str_id): num_elems = int(np.sum([np.prod(x) for x in input_size])) name_template = "{}__batch_size_{}_input_element_size_{}_{}" base_dataset = input_dataset.map(map_fn).batch(batch_size) options = dataset_ops.Options() opt_options = optimization_options.OptimizationOptions() # Disable default map_and_batch_fusion optimization opt_options.map_and_batch_fusion = False options.experimental_optimization = opt_options base_dataset = base_dataset.with_options(options) unoptimized_op = dataset_ops.make_one_shot_iterator( base_dataset).get_next() optimized_options = dataset_ops.Options() opt_options = optimization_options.OptimizationOptions() opt_options.map_vectorization = True optimized_options.experimental_optimization = opt_options optimized = base_dataset.with_options(optimized_options) optimized_op = dataset_ops.make_one_shot_iterator(optimized).get_next() unoptimized_time = self._run(unoptimized_op, name=name_template.format( str_id, batch_size, num_elems, "unoptimized")) optimized_time = self._run(optimized_op, name=name_template.format( str_id, batch_size, num_elems, "optimized")) print("Batch size: {}\n" "Input element size: {}\n" "Transformation: {}\n" "Speedup: {}\n".format(batch_size, input_size, str_id, (unoptimized_time / optimized_time)))
def testSkipEagerOptimizationWithCapturedRefVar(self, dataset_fn): """Tests that default optimizations are disabled with ref variables.""" variable = variable_scope.get_variable("v", initializer=0, use_resource=False) assign_op = variable.assign_add(1) unoptimized_dataset = dataset_fn(variable) options = dataset_ops.Options() opt_options = optimization_options.OptimizationOptions() opt_options.noop_elimination = True opt_options.map_and_batch_fusion = True options.experimental_optimization = opt_options optimized_dataset = unoptimized_dataset.with_options(options) # Check that warning is logged. warnings.simplefilter("always") with warnings.catch_warnings(record=True) as w: optimized_it = optimized_dataset.make_initializable_iterator() self.assertGreaterEqual(len(w), 1) expected = ( "tf.data static optimizations are not compatible with " "tf.Variable. The following optimizations will be disabled: %s." " To enable optimizations, use resource variables instead by " "calling `tf.enable_resource_variables()` at the start of the " "program." % (", ".join(opt_options._static_optimizations()))) self.assertTrue(any([expected in str(warning) for warning in w])) # Check that outputs are the same in the optimized and unoptimized cases, # when the variable value is changing. unoptimized_it = unoptimized_dataset.make_initializable_iterator() with ops.control_dependencies([assign_op]): unoptimized_output = unoptimized_it.get_next() optimized_output = optimized_it.get_next() self.evaluate(variable.initializer) self.evaluate((unoptimized_it.initializer, optimized_it.initializer)) while True: try: unoptimized, optimized = self.evaluate( (unoptimized_output, optimized_output)) self.assertEqual(unoptimized, optimized) except errors.OutOfRangeError: break
def _get_test_datasets(self, base_dataset, map_fn, num_parallel_calls=None, expect_optimized=True): """Given base dataset and map fn, creates test datasets. Returns a tuple of (unoptimized dataset, optimized dataset). The unoptimized dataset has the assertion that Batch follows Map. The optimized dataset has the assertion that Map follows Batch, and has the "map_vectorization" optimization applied. Args: base_dataset: Input dataset to map->batch map_fn: Map function to use num_parallel_calls: (Optional.) num_parallel_calls argument for map expect_optimized: (Optional.) Whether we expect the optimization to take place, in which case we will assert that Batch is followed by Map, otherwise Map followed by Batch. Defaults to True. Returns: Tuple of (unoptimized dataset, optimized dataset). """ map_node_name = "Map" if num_parallel_calls is None else "ParallelMap" def _make_dataset(node_names): dataset = base_dataset.apply(optimization.assert_next(node_names)) dataset = dataset.map(map_fn, num_parallel_calls) dataset = dataset.batch(100) options = dataset_ops.Options() opt_options = optimization_options.OptimizationOptions() opt_options.map_and_batch_fusion = False options.experimental_optimization = opt_options dataset = dataset.with_options(options) return dataset unoptimized = _make_dataset([map_node_name, "Batch"]) optimized = _make_dataset(["Batch", map_node_name] if expect_optimized else [map_node_name, "Batch"]) options = dataset_ops.Options() opt_options = optimization_options.OptimizationOptions() opt_options.map_vectorization = True options.experimental_optimization = opt_options optimized = optimized.with_options(options) return unoptimized, optimized
def testOptimizationNestedDatasetWithModifiedRetval(self): def flat_map_fn(_): dataset = dataset_ops.Dataset.from_tensors(0) dataset = dataset.apply(optimization.assert_next(["MapAndBatch"])) # Should be fused by map and batch fusion dataset = dataset.map(lambda x: x) dataset = dataset.batch(1) return dataset dataset = dataset_ops.Dataset.range(1) dataset = dataset.flat_map(flat_map_fn) # TODO(b/120558523): We use Options instead of _OptimizeDataset directly # here because of a bug with chaining _OptimizeDatasets when there are # nested dataset functions options = dataset_ops.Options() opt_options = optimization_options.OptimizationOptions() opt_options.map_and_batch_fusion = True options.experimental_optimization = opt_options dataset = dataset.with_options(options) self.assertDatasetProduces(dataset, expected_output=[[0]])
def __init__(self, dataset, devices, max_buffer_size=1, prefetch_buffer_size=1, source_device="/cpu:0"): """Constructs a MultiDeviceIterator. Args: dataset: The input dataset to be iterated over. devices: The list of devices to fetch data to. max_buffer_size: Maximum size of the host side per device buffer to keep. prefetch_buffer_size: if > 1, then we setup a buffer on each device to prefetch into. source_device: The host device to place the `dataset` on. Raises: RuntimeError: If run in Eager mode. """ if context.executing_eagerly(): # TODO(rohanj): Fix this. Tracking bug: b/116467184 raise RuntimeError( "MultiDeviceIterator is not currently supported in " "Eager mode.") self._dataset = dataset._apply_options() # pylint: disable=protected-access self._devices = devices self._source_device = source_device self._source_device_tensor = ops.convert_to_tensor(source_device) # Create the MultiDeviceIterator. with ops.device(self._source_device): self._multi_device_iterator_resource = ( gen_dataset_ops.multi_device_iterator( devices=self._devices, shared_name="", container="", **dataset_ops.flat_structure(dataset))) # The incarnation ID is used to ensure consistency between the per-device # iterators and the multi-device iterator. self._incarnation_id = gen_dataset_ops.multi_device_iterator_init( self._dataset._as_variant_tensor(), # pylint: disable=protected-access self._multi_device_iterator_resource, max_buffer_size=max_buffer_size) # TODO(rohanj): Explore the possibility of the MultiDeviceIterator to # initialize the device side of the pipeline. This would allow the # MultiDeviceIterator to choose, for example, to move some transformations # into the device side from its input. It might be useful in rewriting. # Create the per device iterators. self._device_iterators = [] for i, device in enumerate(self._devices): ds = _PerDeviceGenerator(i, self._multi_device_iterator_resource, self._incarnation_id, self._source_device_tensor, device, dataset._element_structure) # pylint: disable=protected-access if prefetch_buffer_size > 0: ds = ds.prefetch(prefetch_buffer_size) # TODO(jsimsa): Enable auto-tuning and optimizations when supported for # non-CPU devices. options = dataset_ops.Options() options.experimental_autotune = False opt_options = optimization_options.OptimizationOptions() opt_options.apply_default_optimizations = False options.experimental_optimization = opt_options ds = ds.with_options(options) with ops.device(device): self._device_iterators.append(ds.make_initializable_iterator()) device_iterator_initializers = [ iterator.initializer for iterator in self._device_iterators ] self._initializer = control_flow_ops.group( *device_iterator_initializers)