def testOptionsProtoRoundTrip(self): options = options_lib.Options() options.autotune.enabled = True options.autotune.cpu_budget = 10 options.autotune.ram_budget = 20 options.deterministic = True options.experimental_external_state_policy = ( options_lib.ExternalStatePolicy.FAIL) options.experimental_distribute.auto_shard_policy = ( options_lib.AutoShardPolicy.DATA) options.experimental_distribute.num_devices = 1000 options.experimental_optimization.apply_default_optimizations = True options.experimental_optimization.filter_fusion = True options.experimental_optimization.map_and_batch_fusion = True options.experimental_optimization.map_and_filter_fusion = True options.experimental_optimization.map_fusion = True options.experimental_optimization.map_parallelization = True options.experimental_optimization.noop_elimination = True options.experimental_optimization.parallel_batch = True options.experimental_optimization.shuffle_and_repeat_fusion = True options.experimental_slack = True options.threading.max_intra_op_parallelism = 30 options.threading.private_threadpool_size = 40 pb = options._to_proto() result = options_lib.Options() result._from_proto(pb) self.assertEqual(options, result)
def testOptionsHaveDefaults(self): options1 = options_lib.Options() options2 = options_lib.Options() self.assertIsNot(options1.experimental_optimization, options2.experimental_optimization) self.assertIsNot(options1.threading, options2.threading) self.assertEqual(options1.experimental_optimization, options_lib.OptimizationOptions()) self.assertEqual(options1.threading, options_lib.ThreadingOptions())
def testExperimentalThreadingOptionsOverride(self): options = options_lib.Options() self.assertEqual(options.threading, options.experimental_threading) options.threading.max_intra_op_parallelism = 20 options.experimental_threading.max_intra_op_parallelism = 40 pb = options._to_proto() result = options_lib.Options() result._from_proto(pb) self.assertEqual(result.experimental_threading.max_intra_op_parallelism, result.threading.max_intra_op_parallelism)
def testOptionsMergeOptionsFromMultipleInputs(self): options1 = options_lib.Options() options1.autotune.enabled = True options2 = options_lib.Options() options2.deterministic = True ds1 = dataset_ops.Dataset.range(0).with_options(options1) ds2 = dataset_ops.Dataset.range(0).with_options(options2) ds = dataset_ops.Dataset.zip((ds1, ds2)) options = self._get_options(ds) self.assertTrue(options.autotune.enabled) self.assertTrue(options.deterministic)
def testExperimentalDeterministicOverride(self): options = options_lib.Options() self.assertEqual(options.deterministic, options.experimental_deterministic) options.experimental_deterministic = False pb = options._to_proto() result = options_lib.Options() result._from_proto(pb) self.assertFalse(result.deterministic) self.assertEqual(result.deterministic, result.experimental_deterministic) result.experimental_deterministic = True self.assertTrue(result.deterministic) self.assertEqual(result.deterministic, result.experimental_deterministic)
def testMutatingOptionsRaiseValueError(self): ds = dataset_ops.Dataset.range(0) options1 = options_lib.Options() options1.experimental_slack = True options2 = options_lib.Options() options2.autotune.enabled = True ds = ds.with_options(options1) ds = ds.map(lambda x: 2 * x) ds = ds.with_options(options2) dataset_options = ds.options() with self.assertRaises(ValueError): dataset_options.deterministic = True
def testOptionsTwiceDifferentOptions(self): options1 = options_lib.Options() options1.autotune.enabled = True options2 = options_lib.Options() options2.deterministic = False ds = dataset_ops.Dataset.range(0) ds = ds.with_options(options1) ds = ds.with_options(options2) options = self._get_options(ds) self.assertTrue(options.autotune.enabled) # Explicitly check that flag is False since assertFalse allows None self.assertIs(options.deterministic, False)
def testOptionsTwiceSameOption(self): if sys.version_info >= (3, 8) and platform.system() == "Windows": # TODO(b/165013260): Fix this self.skipTest("Test is currently broken on Windows with Python 3.8") options1 = options_lib.Options() options1.autotune.enabled = False options2 = options_lib.Options() options2.autotune.enabled = True ds = dataset_ops.Dataset.range(0) ds = ds.with_options(options1) ds = ds.with_options(options2) self.assertTrue(self._get_options(ds).autotune.enabled)
def testOptimizationInjectPrefetch(self, existing_prefetch, autotune, set_env): if set_env: os.environ["TF_DATA_EXPERIMENT_OPT_IN"] = "inject_prefetch" os.environ["TF_JOB_NAME"] = "test_job" dataset = dataset_ops.Dataset.range(5) dataset = dataset.map(lambda x: x + 1, num_parallel_calls=dataset_ops.AUTOTUNE) if existing_prefetch: dataset = dataset.prefetch(1) if autotune and set_env and not existing_prefetch: dataset = dataset.apply(testing.assert_next(["Prefetch", "Root"])) else: dataset = dataset.apply(testing.assert_next(["Root"])) options = options_lib.Options() options.autotune.enabled = autotune dataset = dataset.with_options(options) self.assertDatasetProduces(dataset, expected_output=list(range(1, 6))) if set_env: del os.environ["TF_DATA_EXPERIMENT_OPT_IN"] del os.environ["TF_JOB_NAME"]
def testThreadingOptionsBackwardCompatibility(self): opts = options_lib.Options() opts.threading.max_intra_op_parallelism = 20 self.assertEqual(opts.experimental_threading.max_intra_op_parallelism, 20) opts.experimental_threading.private_threadpool_size = 80 self.assertEqual(opts.threading.private_threadpool_size, 80)
def testStatefulExternalPolicy(self): checkpoint_directory = self.get_temp_dir() checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") dataset = dataset_ops.Dataset.range(4) def fn(x): return x * x dataset = dataset.map( lambda x: script_ops.eager_py_func(fn, [x], dtypes.int64)) options = options_lib.Options() options.experimental_external_state_policy = ( options_lib.ExternalStatePolicy.WARN) dataset = dataset.with_options(options) iterator = iter(dataset) get_next = iterator.get_next checkpoint = trackable_utils.Checkpoint(iterator=iterator) self.assertEqual(0, get_next().numpy()) self.assertEqual(1, get_next().numpy()) save_path = checkpoint.save(checkpoint_prefix) self.assertEqual(4, get_next().numpy()) self.assertEqual(9, get_next().numpy()) checkpoint.restore(save_path).run_restore_ops() self.assertEqual(4, get_next().numpy()) self.assertEqual(9, get_next().numpy()) with self.assertRaises(errors.OutOfRangeError): get_next()
def test_stateful_ops_map_and_batch(self, use_function, use_legacy_map_and_batch): with test_util.deterministic_ops(): v = variables.Variable(0.) def map_fn(x): v.assign_add(1.) return (x, v.read_value()) if use_function: map_fn = def_function.function(map_fn) dataset = dataset_ops.Dataset.range(5) dataset = dataset.apply(testing.assert_next(["MapAndBatch"])) if use_legacy_map_and_batch: dataset = dataset.apply( batching.map_and_batch(map_fn, 2, num_parallel_calls=5)) else: dataset = dataset.map(map_fn, num_parallel_calls=5) dataset = dataset.batch(2) options = options_lib.Options() options.experimental_optimization.apply_default_optimizations = False options.experimental_optimization.map_and_batch_fusion = True dataset = dataset.with_options(options) self.evaluate(variables.global_variables_initializer()) expected_output = [ (np.array([0, 1]), np.array([1, 2])), (np.array([2, 3]), np.array([3, 4])), (np.array([4]), np.array([5])), ] self.assertDatasetProduces(dataset, expected_output=expected_output, requires_initialization=True)
def testDeterminism(self, local_determinism, global_determinism): num_elements = 1000 batches = [] for i in range(num_elements): example_i = example(features=features({ "a": int64_feature([i]), })) batches.append([example_i.SerializeToString()]) test_features = {"a": parsing_ops.FixedLenFeature((), dtype=dtypes.int64)} dataset = dataset_ops.Dataset.from_tensor_slices(batches) dataset = dataset.apply( contrib_parsing_ops.parse_example_dataset( test_features, num_parallel_calls=10, deterministic=local_determinism)) opts = options_lib.Options() opts.deterministic = global_determinism dataset = dataset.with_options(opts) expected = list(range(num_elements)) actual = [elem["a"][0] for elem in self.getDatasetOutput(dataset)] require_order = local_determinism or (local_determinism is None and global_determinism) if require_order: self.assertAllEqual(expected, actual) else: self.assertCountEqual(expected, actual)
def testAssertNext(self): dataset = dataset_ops.Dataset.from_tensors(0).apply( testing.assert_next(["Map"])).map(lambda x: x) options = options_lib.Options() options.experimental_optimization.apply_default_optimizations = False dataset = dataset.with_options(options) self.assertDatasetProduces(dataset, expected_output=[0])
def testNoopElimination(self, init_dataset_fn, transformation, expected_name): """Runs a noop elimination test case. Args: init_dataset_fn: Function to create the initial dataset transformation: Transformation to apply expected_name: Name of the transformation if it is not eliminated """ dataset = init_dataset_fn() if expected_name: dataset = dataset.apply( testing.assert_next([expected_name, "FiniteTake"])) else: dataset = dataset.apply(testing.assert_next(["FiniteTake"])) dataset = dataset.apply(transformation) dataset = dataset.take(1) options = options_lib.Options() options.experimental_optimization.apply_default_optimizations = False options.experimental_optimization.noop_elimination = True dataset = dataset.with_options(options) # Run the first iteration for the side effect of checking the assertion. get_next = self.getNext(dataset) self.evaluate(get_next())
def testExternalStatePolicyFail(self): with ops.device(self._device0): dataset0 = dataset_ops.Dataset.range(100).map( lambda _: random_ops.random_uniform( # pylint:disable=g-long-lambda [], minval=1, maxval=10, dtype=dtypes.float32)) opt = options_lib.Options() opt.experimental_external_state_policy = ( options_lib.ExternalStatePolicy.FAIL) dataset0 = dataset0.with_options(opt) with self.assertRaises(errors.FailedPreconditionError): replicated_ds = distribute.replicate( dataset0, [self._device1, self._device2]) dataset1 = replicated_ds[self._device1] dataset2 = replicated_ds[self._device2] with ops.device(self._device0): get_next0 = self.getNext(dataset0) with ops.device(self._device1): get_next1 = self.getNext(dataset1) with ops.device(self._device2): get_next2 = self.getNext(dataset2) for _ in range(100): self.evaluate(get_next0()) self.evaluate(get_next1()) self.evaluate(get_next2())
def testOptimizationWithCapturedRefVar(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 = options_lib.Options() options.experimental_optimization.apply_default_optimizations = False options.experimental_optimization.noop_elimination = True options.experimental_optimization.map_and_batch_fusion = True optimized_dataset = unoptimized_dataset.with_options(options) optimized_it = dataset_ops.make_initializable_iterator( optimized_dataset) # Check that outputs are the same in the optimized and unoptimized cases, # when the variable value is changing. unoptimized_it = dataset_ops.make_initializable_iterator( unoptimized_dataset) 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 _benchmark_filter_parallelization(self, chain_length, optimize_dataset): dataset = dataset_ops.Dataset.from_tensors(5).repeat() for _ in range(chain_length): dataset = dataset.filter( lambda x: math_ops.greater_equal(x - 5, 0)) if optimize_dataset: options = options_lib.Options() options.experimental_optimization.apply_default_optimizations = False options.experimental_optimization.filter_parallelization = True dataset = dataset.with_options(options) opt_mark = "opt" if optimize_dataset else "noopt" self.run_and_report_benchmark( dataset=dataset, num_elements=100, iters=10, warmup=True, extras={ "model_name": "optimize.benchmark.4", "parameters": "%d.%s" % (chain_length, optimize_dataset), }, name="filter_parallelization_{}_chain_length_{}".format( opt_mark, chain_length))
def test_deterministic_attribute(self, local_determinism, global_determinism): self._set_seed() with test_util.deterministic_ops(): def sleep(x): time.sleep(0.1) return x def map_function(x): if math_ops.equal(x, 0): return script_ops.py_func(sleep, [x], x.dtype, stateful=False) else: return x dataset = dataset_ops.Dataset.range(100) dataset = dataset.map(map_function, num_parallel_calls=2, deterministic=local_determinism) opts = options_lib.Options() opts.deterministic = global_determinism dataset = dataset.with_options(opts) self.assertDatasetProduces(dataset, expected_output=range(100))
def benchmark_map_and_batch(self): """Measures the performance of parallelized batching.""" shapes = [(), (10, ), (10, 10), (10, 10, 10), (224, 224, 3)] batch_size_values = [1, 32, 64, 128, 1024] for shape in shapes: for batch_size in batch_size_values: dataset = dataset_ops.Dataset.range(1000000000) dense_value = random_ops.random_normal(shape=shape) dataset = dataset.apply( batching.map_and_batch(lambda _: dense_value, batch_size)) # pylint: disable=cell-var-from-loop options = options_lib.Options() options.experimental_optimization.apply_default_optimizations = False dataset = dataset.with_options(options) self.run_and_report_benchmark( dataset=dataset, num_elements=batch_size, iters=100, warmup=True, extras={ "model_name": "map_and_batch.benchmark.1", "parameters": "%d.%s" % (batch_size, str(shape)) }, name="num_elements_%d_batch_size_%d" % (np.prod(shape), batch_size))
def testReplicateAndShardProduceDisjointData(self, shuffle, sharding_policy): dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=shuffle) dataset = dataset.flat_map(core_readers.TFRecordDataset) graph_def = dataset._as_serialized_graph( strip_device_assignment=True, external_state_policy=options_lib.ExternalStatePolicy.WARN) options = options_lib.Options() options.experimental_distribute.auto_shard_policy = sharding_policy ds1 = distribute._RemoteDataset(graph_def, "/device:CPU:0", dataset.element_spec) ds2 = distribute._RemoteDataset(graph_def, "/device:CPU:0", dataset.element_spec) ds1 = ds1.with_options(options) ds2 = ds2.with_options(options) ds1 = distribute._AutoShardDataset(ds1, 2, 0) ds2 = distribute._AutoShardDataset(ds2, 2, 1) elems1 = set(self.getAllDatasetElements(ds1)) elems2 = set(self.getAllDatasetElements(ds2)) self.assertEmpty(elems1.intersection(elems2))
def test_text_line_dataset(self, use_function): self._set_seed() with test_util.deterministic_ops(): def write_nums_to_file(filename, numbers): path = os.path.join(self.get_temp_dir(), filename) with open(path, "w") as f: f.write("\n".join(str(n) for n in numbers)) return path f1 = write_nums_to_file("f1", (1, 2, 3)) f2 = write_nums_to_file("f2", (4, 5, 6)) f3 = write_nums_to_file("f3", (7, 8, 9)) def interleave_fn(filename): return reader_ops.TextLineDataset(filename) if use_function: interleave_fn = def_function.function(interleave_fn) dataset = dataset_ops.Dataset.from_tensor_slices([f1, f2, f3]) dataset = dataset.apply(testing.assert_next(["ParallelInterleave" ])) dataset = dataset.interleave(interleave_fn, cycle_length=3, num_parallel_calls=3) options = options_lib.Options() options.experimental_optimization.apply_default_optimizations = False dataset = dataset.with_options(options) self.assertDatasetProduces( dataset, expected_output=["1", "4", "7", "2", "5", "8", "3", "6", "9"])
def test_stateful_ops_interleave(self, use_function, use_legacy_interleave): with test_util.deterministic_ops(): v = variables.Variable(0.) def map_fn(x): v.assign_add(1.) return (x, v.read_value()) def interleave_fn(x): del x return dataset_ops.Dataset.range(2).map(map_fn) if use_function: map_fn = def_function.function(map_fn) interleave_fn = def_function.function(interleave_fn) dataset = dataset_ops.Dataset.range(5) if use_legacy_interleave: dataset = dataset.apply( interleave_ops.parallel_interleave(interleave_fn, cycle_length=5)) else: dataset = dataset.interleave(interleave_fn, cycle_length=5, num_parallel_calls=3) options = options_lib.Options() options.experimental_optimization.apply_default_optimizations = False dataset = dataset.with_options(options) self.evaluate(variables.global_variables_initializer()) expected_output = list(zip([0] * 5 + [1] * 5, range(1, 11))) self.assertDatasetProduces(dataset, expected_output=expected_output, requires_initialization=True)
def testFilterFusion(self, function, predicates): dataset = dataset_ops.Dataset.range(5).apply( testing.assert_next(["Map", "Filter", "MemoryCacheImpl"])).map(function) for predicate in predicates: dataset = dataset.filter(predicate) dataset = dataset.cache() options = options_lib.Options() options.experimental_optimization.apply_default_optimizations = False options.experimental_optimization.filter_fusion = True dataset = dataset.with_options(options) expected_output = [] for x in range(5): r = function(x) filtered = False for predicate in predicates: if isinstance(r, tuple): b = predicate(*r) # Pass tuple as multiple arguments. else: b = predicate(r) if not self.evaluate(b): filtered = True break if not filtered: expected_output.append(r) self.assertDatasetProduces(dataset, expected_output=expected_output)
def test_no_stateful_ops_interleave(self, use_function, use_legacy_interleave): self._set_seed() with test_util.deterministic_ops(): def interleave_fn(x): del x return dataset_ops.Dataset.range(2) if use_function: interleave_fn = def_function.function(interleave_fn) dataset = dataset_ops.Dataset.range(5) if use_legacy_interleave: dataset = dataset.apply( testing.assert_next(["LegacyParallelInterleaveV2"])) dataset = dataset.apply( interleave_ops.parallel_interleave(interleave_fn, cycle_length=5)) else: dataset = dataset.apply( testing.assert_next(["ParallelInterleave"])) dataset = dataset.interleave(interleave_fn, cycle_length=5, num_parallel_calls=3) options = options_lib.Options() options.experimental_optimization.apply_default_optimizations = False dataset = dataset.with_options(options) self.evaluate(variables.global_variables_initializer()) self.assertDatasetProduces(dataset, expected_output=[0] * 5 + [1] * 5)
def benchmark_resample_performance(self): init_dist = [0.25, 0.25, 0.25, 0.25] target_dist = [0.0, 0.0, 0.0, 1.0] num_classes = len(init_dist) # We don't need many samples to test a dirac-delta target distribution num_samples = 1000 data_np = np.random.choice(num_classes, num_samples, p=init_dist) # Prepare the dataset dataset = dataset_ops.Dataset.from_tensor_slices(data_np).repeat() # Reshape distribution via rejection sampling. dataset = dataset.apply( resampling.rejection_resample(class_func=lambda x: x, target_dist=target_dist, initial_dist=init_dist, seed=142)) options = options_lib.Options() options.experimental_optimization.apply_default_optimizations = False dataset = dataset.with_options(options) wall_time = self.run_benchmark(dataset=dataset, num_elements=num_samples, iters=10, warmup=True) resample_time = wall_time * num_samples self.report_benchmark(iters=10, wall_time=resample_time, extras={ "model_name": "rejection_resample.benchmark.1", "parameters": "%d" % num_samples, }, name="resample_{}".format(num_samples))
def testNoErrorWithoutPrefetch(self): """The rewrite should not fail if there is no prefetch() in the pipeline.""" dataset = dataset_ops.Dataset.range(10) options = options_lib.Options() options.experimental_slack = True dataset = dataset.with_options(options) self.assertDatasetProduces(dataset, range(10))
def run_benchmark(self, dataset, num_elements, iters=1, warmup=True, apply_default_optimizations=False, session_config=None): """Benchmarks the dataset. Runs the dataset `iters` times. In each iteration, the benchmark measures the time it takes to go through `num_elements` elements of the dataset. Args: dataset: Dataset to benchmark. num_elements: Number of dataset elements to iterate through each benchmark iteration. iters: Number of times to repeat the timing. warmup: If true, warms up the session caches by running an untimed run. apply_default_optimizations: Determines whether default optimizations should be applied. session_config: A ConfigProto protocol buffer with configuration options for the session. Applicable only for benchmarking in graph mode. Returns: A float, representing the per-element wall time of the dataset in seconds. This is the median time (with respect to `iters`) it takes for the dataset to go through `num_elements` elements, divided by `num_elements.` """ # The options that have been applied to the dataset are preserved so that # they are not overwritten while benchmarking. options = options_lib.Options() options.experimental_optimization.apply_default_optimizations = ( apply_default_optimizations) dataset = dataset.with_options(options) # NOTE: We use `dataset.skip()` to perform the iterations in C++, avoiding # the overhead of having to execute a TensorFlow op for each step of the # input pipeline. Note that this relies on the underlying implementation of # `skip` to execute upstream computation. If it is optimized in the future, # we will have to change this code. dataset = dataset.skip(num_elements - 1) if context.executing_eagerly(): median_duration = self._run_eager_benchmark(iterable=dataset, iters=iters, warmup=warmup) return median_duration / float(num_elements) iterator = dataset_ops.make_initializable_iterator(dataset) next_element = iterator.get_next() op = nest.flatten(next_element)[0].op median_duration = self._run_graph_benchmark( iterable=op, iters=iters, warmup=warmup, session_config=session_config, initializer=iterator.initializer) return median_duration / float(num_elements)
def testPrefetchWithSlackOptionWithoutIterator(self): """Defaults to slack period of 1 without iterator.""" dataset = dataset_ops.Dataset.range(10) dataset = dataset.prefetch(1) options = options_lib.Options() options.experimental_slack = True dataset = dataset.with_options(options) self.assertDatasetProduces(dataset, range(10))
def testOptimizationStatefulFunction(self): dataset = dataset_ops.Dataset.range(10).map( lambda _: random_ops.random_uniform([])).batch(10) options = options_lib.Options() options.experimental_optimization.apply_default_optimizations = False dataset = dataset.with_options(options) get_next = self.getNext(dataset) self.evaluate(get_next())