chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
License Check / license-check (push) Failing after 2s

This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
@@ -0,0 +1,150 @@
load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
load(
"//tensorflow/tools/test:performance.bzl",
"tf_py_benchmark_test",
)
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//tensorflow:internal"],
licenses = ["notice"],
)
py_binary(
name = "autotune_benchmark_binary",
srcs = ["autotune_benchmark.py"],
main = "autotune_benchmark.py",
strict_deps = True,
deps = [
"//tensorflow/python/data/benchmarks:benchmark_base",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:options",
"//tensorflow/python/ops:math_ops",
"//third_party/py/numpy",
],
)
tf_py_benchmark_test(
name = "autotune_benchmark",
srcs = ["autotune_benchmark.py"],
deps = [
"//tensorflow/python/data/benchmarks:benchmark_base",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:options",
"//tensorflow/python/ops:math_ops",
"//third_party/py/numpy",
],
)
tf_py_benchmark_test(
name = "parameter_value_benchmark",
srcs = ["parameter_value_benchmark.py"],
deps = [
"//tensorflow/python/data/benchmarks:benchmark_base",
"//tensorflow/python/data/experimental/ops:testing",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:script_ops",
"//third_party/py/numpy",
],
)
tf_py_benchmark_test(
name = "csv_dataset_benchmark",
srcs = ["csv_dataset_benchmark.py"],
tags = ["no_pip"],
deps = [
"//tensorflow/python/data/benchmarks:benchmark_base",
"//tensorflow/python/data/experimental/ops:readers",
"//tensorflow/python/data/ops:readers",
"//tensorflow/python/ops:parsing_ops",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/platform:test",
],
)
tf_py_benchmark_test(
name = "map_and_batch_benchmark",
srcs = ["map_and_batch_benchmark.py"],
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/data/benchmarks:benchmark_base",
"//tensorflow/python/data/experimental/ops:batching",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:options",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:random_ops",
"//third_party/py/numpy",
],
)
tf_py_benchmark_test(
name = "map_defun_benchmark",
srcs = ["map_defun_benchmark.py"],
deps = [
"//tensorflow/python/data/benchmarks:benchmark_base",
"//tensorflow/python/data/experimental/ops:map_defun",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:map_fn",
"//tensorflow/python/ops:math_ops",
],
)
tf_py_benchmark_test(
name = "matching_files_benchmark",
size = "small",
srcs = ["matching_files_benchmark.py"],
deps = [
"//tensorflow/python/data/benchmarks:benchmark_base",
"//tensorflow/python/data/experimental/ops:matching_files",
],
)
tf_py_benchmark_test(
name = "optimize_benchmark",
srcs = ["optimize_benchmark.py"],
deps = [
"//tensorflow/python/data/benchmarks:benchmark_base",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:options",
"//tensorflow/python/ops:math_ops",
],
)
tf_py_benchmark_test(
name = "rejection_resample_benchmark",
srcs = ["rejection_resample_benchmark.py"],
tags = ["no_pip"],
deps = [
"//tensorflow/python/data/benchmarks:benchmark_base",
"//tensorflow/python/data/experimental/ops:resampling",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:options",
"//third_party/py/numpy",
],
)
tf_py_benchmark_test(
name = "snapshot_dataset_benchmark",
srcs = ["snapshot_dataset_benchmark.py"],
deps = [
"//tensorflow/python/data/benchmarks:benchmark_base",
"//tensorflow/python/data/experimental/ops:snapshot",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/ops:array_ops_gen",
"//tensorflow/python/platform:client_testlib",
],
)
tf_py_benchmark_test(
name = "unbatch_benchmark",
srcs = ["unbatch_benchmark.py"],
deps = [
"//tensorflow/python/data/benchmarks:benchmark_base",
"//tensorflow/python/data/ops:dataset_ops",
],
)
@@ -0,0 +1,216 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Benchmarks for autotuning performance knobs."""
import numpy as np
from tensorflow.python.data.benchmarks import benchmark_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import options as options_lib
from tensorflow.python.ops import math_ops
class AutotuneBenchmark(benchmark_base.DatasetBenchmarkBase):
"""Benchmarks for autotuning performance knobs."""
def _run_benchmark(self, dataset, autotune, benchmark_iters, benchmark_label,
benchmark_id):
options = options_lib.Options()
options.experimental_optimization.apply_default_optimizations = False
options.autotune.enabled = autotune
dataset = dataset.with_options(options)
autotune_string = "_autotune_parallelism_only"
wall_time = self.run_and_report_benchmark(
dataset=dataset,
num_elements=benchmark_iters,
warmup=True,
iters=1,
extras={
"model_name":
"autotune.benchmark.%s.%d" % (benchmark_label, benchmark_id),
"parameters":
"%s" % autotune,
},
name=benchmark_label + (autotune_string if autotune else ""))
return wall_time
def benchmark_batch(self):
a = self._benchmark_batch(autotune=False, benchmark_id=1)
b = self._benchmark_batch(autotune=True, benchmark_id=2)
print("autotune parallelism vs no autotuning speedup: {}".format(a / b))
def _benchmark_batch(self, autotune, benchmark_id):
batch_size = 128
k = 1024
dataset = dataset_ops.Dataset.from_tensors(
(np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))).repeat()
dataset = dataset.map(math_ops.matmul)
dataset = dataset.batch(
batch_size=batch_size, num_parallel_calls=dataset_ops.AUTOTUNE)
return self._run_benchmark(
dataset=dataset,
autotune=autotune,
benchmark_iters=10000,
benchmark_label="batch",
benchmark_id=benchmark_id)
def benchmark_map(self):
a = self._benchmark_map(autotune=False, benchmark_id=1)
b = self._benchmark_map(autotune=True, benchmark_id=2)
print("autotune parallelism vs no autotuning speedup: {}".format(a / b))
def _benchmark_map(self, autotune, benchmark_id):
k = 1024 * 1024
dataset = dataset_ops.Dataset.from_tensors(
(np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))).repeat()
dataset = dataset.map(
math_ops.matmul, num_parallel_calls=dataset_ops.AUTOTUNE)
return self._run_benchmark(
dataset=dataset,
autotune=autotune,
benchmark_iters=10000,
benchmark_label="map",
benchmark_id=benchmark_id)
def benchmark_map_and_batch(self):
a = self._benchmark_map_and_batch(autotune=False, benchmark_id=1)
b = self._benchmark_map_and_batch(autotune=True, benchmark_id=2)
print("autotune parallelism vs no autotuning speedup: {}".format(a / b))
def _benchmark_map_and_batch(self, autotune, benchmark_id):
batch_size = 16
k = 1024 * 1024
dataset = dataset_ops.Dataset.from_tensors(
(np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))).repeat()
dataset = dataset.map(
math_ops.matmul, num_parallel_calls=dataset_ops.AUTOTUNE)
dataset = dataset.batch(batch_size=batch_size)
return self._run_benchmark(
dataset=dataset,
autotune=autotune,
benchmark_iters=1000,
benchmark_label="map_and_batch",
benchmark_id=benchmark_id)
def benchmark_interleave(self):
a = self._benchmark_interleave(autotune=False, benchmark_id=1)
b = self._benchmark_interleave(autotune=True, benchmark_id=2)
print("autotune parallelism vs no autotuning speedup: {}".format(a / b))
def _benchmark_interleave(self, autotune, benchmark_id):
k = 1024 * 1024
dataset = dataset_ops.Dataset.from_tensors(
(np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))).repeat()
dataset = dataset.map(math_ops.matmul)
dataset = dataset_ops.Dataset.range(1).repeat().interleave(
lambda _: dataset,
cycle_length=10,
num_parallel_calls=dataset_ops.AUTOTUNE)
return self._run_benchmark(
dataset=dataset,
autotune=autotune,
benchmark_iters=10000,
benchmark_label="interleave",
benchmark_id=benchmark_id)
def benchmark_map_and_interleave(self):
a = self._benchmark_map_and_interleave(autotune=False, benchmark_id=1)
b = self._benchmark_map_and_interleave(autotune=True, benchmark_id=2)
print("autotune parallelism vs no autotuning speedup: {}".format(a / b))
def _benchmark_map_and_interleave(self, autotune, benchmark_id):
k = 1024 * 1024
a = (np.random.rand(1, 8 * k), np.random.rand(8 * k, 1))
b = (np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))
c = (np.random.rand(1, 2 * k), np.random.rand(2 * k, 1))
dataset_a = dataset_ops.Dataset.from_tensors(a).repeat()
dataset_b = dataset_ops.Dataset.from_tensors(b).repeat()
dataset_c = dataset_ops.Dataset.from_tensors(c).repeat()
def f1(x, y):
return math_ops.matmul(x, y)
def f2(a, b):
x, y = b
return a, math_ops.matmul(x, y)
dataset = dataset_a
dataset = dataset.map(f1, num_parallel_calls=dataset_ops.AUTOTUNE)
dataset = dataset_ops.Dataset.range(1).repeat().interleave(
lambda _: dataset,
num_parallel_calls=dataset_ops.AUTOTUNE,
cycle_length=2)
dataset = dataset_ops.Dataset.zip((dataset, dataset_b))
dataset = dataset.map(f2, num_parallel_calls=dataset_ops.AUTOTUNE)
dataset = dataset_ops.Dataset.range(1).repeat().interleave(
lambda _: dataset,
num_parallel_calls=dataset_ops.AUTOTUNE,
cycle_length=2)
dataset = dataset_ops.Dataset.zip((dataset, dataset_c))
dataset = dataset.map(f2, num_parallel_calls=dataset_ops.AUTOTUNE)
return self._run_benchmark(
dataset=dataset,
autotune=autotune,
benchmark_iters=10000,
benchmark_label="map_and_interleave",
benchmark_id=benchmark_id)
def benchmark_map_batch_and_interleave(self):
a = self._benchmark_map_batch_and_interleave(autotune=False, benchmark_id=1)
b = self._benchmark_map_batch_and_interleave(autotune=True, benchmark_id=2)
print("autotune parallelism vs no autotuning speedup: {}".format(a / b))
def _benchmark_map_batch_and_interleave(self, autotune, benchmark_id):
batch_size = 16
k = 1024 * 1024
a = (np.random.rand(1, 8 * k), np.random.rand(8 * k, 1))
b = (np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))
c = (np.random.rand(1, 2 * k), np.random.rand(2 * k, 1))
dataset_a = dataset_ops.Dataset.from_tensors(a).repeat()
dataset_b = dataset_ops.Dataset.from_tensors(b).repeat()
dataset_c = dataset_ops.Dataset.from_tensors(c).repeat()
dataset = dataset_a
dataset = dataset.map(
math_ops.matmul, num_parallel_calls=dataset_ops.AUTOTUNE)
dataset = dataset.batch(batch_size=batch_size)
dataset = dataset_ops.Dataset.range(1).repeat().interleave(
lambda _: dataset,
num_parallel_calls=dataset_ops.AUTOTUNE,
cycle_length=2)
dataset = dataset_ops.Dataset.zip((dataset, dataset_b))
dataset = dataset_ops.Dataset.range(1).repeat().interleave(
lambda _: dataset,
num_parallel_calls=dataset_ops.AUTOTUNE,
cycle_length=2)
dataset_c = dataset_c.map(
math_ops.matmul, num_parallel_calls=dataset_ops.AUTOTUNE)
dataset_c = dataset_c.batch(batch_size=batch_size)
dataset = dataset_ops.Dataset.zip((dataset, dataset_c))
return self._run_benchmark(
dataset=dataset,
autotune=autotune,
benchmark_iters=1000,
benchmark_label="map_and_batch_and_interleave",
benchmark_id=benchmark_id)
if __name__ == "__main__":
benchmark_base.test.main()
@@ -0,0 +1,126 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Benchmarks for `tf.data.experimental.CsvDataset`."""
import os
import string
import tempfile
from tensorflow.python.data.benchmarks import benchmark_base
from tensorflow.python.data.experimental.ops import readers
from tensorflow.python.data.ops import readers as core_readers
from tensorflow.python.ops import parsing_ops
from tensorflow.python.platform import gfile
from tensorflow.python.platform import googletest
class CsvDatasetBenchmark(benchmark_base.DatasetBenchmarkBase):
"""Benchmarks for `tf.data.experimental.CsvDataset`."""
FLOAT_VAL = '1.23456E12'
STR_VAL = string.ascii_letters * 10
def _set_up(self, str_val):
# Since this isn't test.TestCase, have to manually create a test dir
gfile.MakeDirs(googletest.GetTempDir())
self._temp_dir = tempfile.mkdtemp(dir=googletest.GetTempDir())
self._num_cols = [4, 64, 256]
self._num_per_iter = 5000
self._filenames = []
for n in self._num_cols:
fn = os.path.join(self._temp_dir, 'file%d.csv' % n)
with open(fn, 'w') as f:
# Just write 100 rows and use `repeat`... Assumes the cost
# of creating an iterator is not significant
row = ','.join(str_val for _ in range(n))
f.write('\n'.join(row for _ in range(100)))
self._filenames.append(fn)
def _tear_down(self):
gfile.DeleteRecursively(self._temp_dir)
def _run_benchmark(self, dataset, num_cols, prefix, benchmark_id):
self.run_and_report_benchmark(
dataset=dataset,
num_elements=self._num_per_iter,
name='%s_with_cols_%d' % (prefix, num_cols),
iters=10,
extras={
'model_name': 'csv.benchmark.%d' % benchmark_id,
'parameters': '%d' % num_cols,
},
warmup=True)
def benchmark_map_with_floats(self):
self._set_up(self.FLOAT_VAL)
for i in range(len(self._filenames)):
num_cols = self._num_cols[i]
kwargs = {'record_defaults': [[0.0]] * num_cols}
dataset = core_readers.TextLineDataset(self._filenames[i]).repeat()
dataset = dataset.map(lambda l: parsing_ops.decode_csv(l, **kwargs)) # pylint: disable=cell-var-from-loop
self._run_benchmark(
dataset=dataset,
num_cols=num_cols,
prefix='csv_float_map_decode_csv',
benchmark_id=1)
self._tear_down()
def benchmark_map_with_strings(self):
self._set_up(self.STR_VAL)
for i in range(len(self._filenames)):
num_cols = self._num_cols[i]
kwargs = {'record_defaults': [['']] * num_cols}
dataset = core_readers.TextLineDataset(self._filenames[i]).repeat()
dataset = dataset.map(lambda l: parsing_ops.decode_csv(l, **kwargs)) # pylint: disable=cell-var-from-loop
self._run_benchmark(
dataset=dataset,
num_cols=num_cols,
prefix='csv_strings_map_decode_csv',
benchmark_id=2)
self._tear_down()
def benchmark_csv_dataset_with_floats(self):
self._set_up(self.FLOAT_VAL)
for i in range(len(self._filenames)):
num_cols = self._num_cols[i]
kwargs = {'record_defaults': [[0.0]] * num_cols}
dataset = core_readers.TextLineDataset(self._filenames[i]).repeat()
dataset = readers.CsvDataset(self._filenames[i], **kwargs).repeat() # pylint: disable=cell-var-from-loop
self._run_benchmark(
dataset=dataset,
num_cols=num_cols,
prefix='csv_float_fused_dataset',
benchmark_id=3)
self._tear_down()
def benchmark_csv_dataset_with_strings(self):
self._set_up(self.STR_VAL)
for i in range(len(self._filenames)):
num_cols = self._num_cols[i]
kwargs = {'record_defaults': [['']] * num_cols}
dataset = core_readers.TextLineDataset(self._filenames[i]).repeat()
dataset = readers.CsvDataset(self._filenames[i], **kwargs).repeat() # pylint: disable=cell-var-from-loop
self._run_benchmark(
dataset=dataset,
num_cols=num_cols,
prefix='csv_strings_fused_dataset',
benchmark_id=4)
self._tear_down()
if __name__ == '__main__':
benchmark_base.test.main()
@@ -0,0 +1,199 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Benchmarks for `tf.data.experimental.map_and_batch()`."""
import hashlib
import itertools
import numpy as np
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.data.benchmarks import benchmark_base
from tensorflow.python.data.experimental.ops import batching
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import options as options_lib
from tensorflow.python.framework import constant_op
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
_NUMPY_RANDOM_SEED = 42
class MapAndBatchBenchmark(benchmark_base.DatasetBenchmarkBase):
"""Benchmarks for `tf.data.experimental.map_and_batch()`."""
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 _benchmark_series(self, label, series, benchmark_id):
"""Runs benchmark the given series."""
# Decides a proper number of iterations according to the inputs.
def compute_num_iters(map_num_calls, inter_op, element_size, batch_size):
return 1024 // (
(element_size * batch_size) //
min(12 if map_num_calls == dataset_ops.AUTOTUNE else map_num_calls,
inter_op))
# Makes the dataset based on the inputs.
def make_dataset(map_num_calls, element_size, batch_size, batch_num_calls,
apply_fusion):
k = 1024 * 1024
x = constant_op.constant(np.random.rand(element_size, 4 * k))
y = constant_op.constant(np.random.rand(4 * k, 1))
dataset = dataset_ops.Dataset.range(1000000000000).map(lambda _: (x, y))
dataset = dataset.map(math_ops.matmul, num_parallel_calls=map_num_calls)
dataset = dataset.batch(
batch_size=batch_size, num_parallel_calls=batch_num_calls)
options = options_lib.Options()
options.experimental_optimization.apply_default_optimizations = False
options.experimental_optimization.map_and_batch_fusion = apply_fusion
dataset = dataset.with_options(options)
return dataset
# Makes the name of the dataset based on the inputs.
def make_name(label, map_num_calls, inter_op, element_size, batch_size,
batch_num_calls, apply_fusion):
map_num_calls_str = ("autotuned" if map_num_calls == dataset_ops.AUTOTUNE
else str(map_num_calls))
batch_num_calls_str = (
"autotuned" if batch_num_calls == dataset_ops.AUTOTUNE else
str(1 if batch_num_calls is None else batch_num_calls))
name_str = ("%s_id_%s_map_num_calls_%s_batch_num_calls_%s_inter_op_%d"
"_elem_size_%d_batch_size_%d")
name = (
name_str % (
"fused" if apply_fusion else "chained",
hashlib.sha1((label).encode("utf-8")).hexdigest()[:8],
map_num_calls_str,
batch_num_calls_str,
inter_op,
element_size,
batch_size,
))
return name
for (map_num_calls, inter_op, element_size, batch_size, batch_num_calls,
apply_fusion) in series:
num_iters = compute_num_iters(map_num_calls, inter_op, element_size,
batch_size)
dataset = make_dataset(map_num_calls, element_size, batch_size,
batch_num_calls, apply_fusion)
name = make_name(label, map_num_calls, inter_op, element_size, batch_size,
batch_num_calls, apply_fusion)
session_config = config_pb2.ConfigProto(
inter_op_parallelism_threads=inter_op, use_per_session_threads=True)
self.run_and_report_benchmark(
dataset=dataset,
iters=num_iters,
num_elements=batch_size,
warmup=True,
extras={
"model_name":
"map_and_batch.benchmark.%d" % benchmark_id,
"parameters":
"%d.%d.%d.%d.%d.%s" %
(map_num_calls, inter_op, element_size, batch_size,
batch_num_calls, apply_fusion),
},
session_config=session_config,
name=name)
def benchmark_map_and_batch_chaining_versus_fusing(self):
"""Compares the performance of chaining and fusing map and batch.
NOTE: It is recommended to build the benchmark with
`-c opt --copt=-mavx --copt=-mavx2 --copt=-mfma --copt=-gmlt`
and execute it on a machine with at least 32 CPU cores.
"""
# Sequential pipeline configurations.
seq_elem_size_series = itertools.product([1], [1], [1, 2, 4, 8], [16],
[None], [False, True])
seq_batch_size_series = itertools.product([1], [1], [1], [8, 16, 32, 64],
[None], [False, True])
# Parallel pipeline configuration.
par_elem_size_series = itertools.product([32], [32], [1, 2, 4, 8], [256],
[None], [False, True])
par_batch_size_series = itertools.product([32], [32], [1],
[128, 256, 512, 1024], [None],
[False, True])
par_map_num_calls_series = itertools.product([8, 16, 32, 64], [32], [1],
[512], [None], [False, True])
par_inter_op_series = itertools.product([32], [8, 16, 32, 64], [1], [512],
[None], [False, True])
# Autotuned pipeline configuration.
fused_versus_chained_series = [
(dataset_ops.AUTOTUNE, 32, 1, 16, dataset_ops.AUTOTUNE, False),
(dataset_ops.AUTOTUNE, 32, 1, 16, None, True)
]
np.random.seed(_NUMPY_RANDOM_SEED)
self._benchmark_series(
"Sequential element size evaluation",
seq_elem_size_series,
benchmark_id=2)
self._benchmark_series(
"Sequential batch size evaluation",
seq_batch_size_series,
benchmark_id=3)
self._benchmark_series(
"Parallel element size evaluation",
par_elem_size_series,
benchmark_id=4)
self._benchmark_series(
"Parallel batch size evaluation", par_batch_size_series, benchmark_id=5)
self._benchmark_series(
"Transformation parallelism evaluation",
par_map_num_calls_series,
benchmark_id=6)
self._benchmark_series(
"Threadpool size evaluation", par_inter_op_series, benchmark_id=7)
self._benchmark_series(
"Autotune chained versus fused evaluation",
fused_versus_chained_series,
benchmark_id=8)
if __name__ == "__main__":
benchmark_base.test.main()
@@ -0,0 +1,75 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Benchmarks for MapDefunOp."""
from tensorflow.python.data.benchmarks import benchmark_base
from tensorflow.python.data.experimental.ops import map_defun
from tensorflow.python.eager import def_function
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import map_fn
from tensorflow.python.ops import math_ops
class MapDefunBenchmark(benchmark_base.DatasetBenchmarkBase):
"""Benchmarks for MapDefunOp."""
def _run(self, op, name, num_iters, benchmark_id):
wall_time = self.run_op_benchmark(op=op, iters=num_iters, warmup=True)
zero_division_delta = 1e-100
wall_time = wall_time + zero_division_delta
self.report_benchmark(
name=name,
iters=num_iters,
wall_time=wall_time,
extras={
"examples_per_sec": 1 / float(wall_time),
"model_name": "map_defun.benchmark.%d" % benchmark_id,
"parameters": "%d" % num_iters,
})
def benchmark_defun_vs_map_fn(self):
"""Benchmarks to compare the performance of MapDefun vs tf.map_fn."""
@def_function.function(
input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
def defun(x):
return array_ops.identity(x)
def fn(x):
return array_ops.identity(x)
base = math_ops.range(10000)
for input_size in [10, 100, 1000, 10000]:
num_iters = 10000 // input_size
map_defun_op = map_defun.map_defun(defun, [base], [dtypes.int32], [()])
map_fn_op = map_fn.map_fn(fn, base)
self._run(
op=map_defun_op,
name="with_defun_size_%d" % input_size,
num_iters=num_iters,
benchmark_id=1)
self._run(
op=map_fn_op,
name="without_defun_size_%d" % input_size,
num_iters=num_iters,
benchmark_id=2)
if __name__ == "__main__":
benchmark_base.test.main()
@@ -0,0 +1,67 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Benchmark for the experimental `MatchingFilesDataset`."""
import os
import shutil
import tempfile
from tensorflow.python.data.benchmarks import benchmark_base
from tensorflow.python.data.experimental.ops import matching_files
class MatchingFilesBenchmark(benchmark_base.DatasetBenchmarkBase):
"""Benchmark for the experimental `MatchingFilesDataset`."""
def benchmark_nested_directories(self):
tmp_dir = tempfile.mkdtemp()
width = 500
depth = 10
for i in range(width):
for j in range(depth):
new_base = os.path.join(tmp_dir, str(i),
*[str(dir_name) for dir_name in range(j)])
os.makedirs(new_base)
child_files = ['a.py', 'b.pyc'] if j < depth - 1 else ['c.txt', 'd.log']
for f in child_files:
filename = os.path.join(new_base, f)
open(filename, 'w').close()
patterns = [
os.path.join(tmp_dir, os.path.join(*['**'
for _ in range(depth)]), suffix)
for suffix in ['*.txt', '*.log']
]
# the num_elements depends on the pattern that has been defined above.
# In the current scenario, the num of files are selected based on the
# ['*.txt', '*.log'] patterns. Since the files which match either of these
# patterns are created once per `width`. The num_elements would be:
num_elements = width * 2
dataset = matching_files.MatchingFilesDataset(patterns)
self.run_and_report_benchmark(
dataset=dataset,
iters=3,
num_elements=num_elements,
extras={
'model_name': 'matching_files.benchmark.1',
'parameters': '%d.%d' % (width, depth),
},
name='nested_directory(%d*%d)' % (width, depth))
shutil.rmtree(tmp_dir, ignore_errors=True)
if __name__ == '__main__':
benchmark_base.test.main()
@@ -0,0 +1,165 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Benchmarks for static optimizations."""
from tensorflow.python.data.benchmarks import benchmark_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import options as options_lib
from tensorflow.python.ops import math_ops
class OptimizationBenchmark(benchmark_base.DatasetBenchmarkBase):
"""Benchmarks for static optimizations."""
def benchmark_map_fusion(self):
"""Evaluates performance map of fusion."""
chain_lengths = [0, 1, 2, 5, 10, 20, 50]
for chain_length in chain_lengths:
self._benchmark_map_fusion(
chain_length=chain_length, optimize_dataset=False)
self._benchmark_map_fusion(
chain_length=chain_length, optimize_dataset=True)
def _benchmark_map_fusion(self, chain_length, optimize_dataset):
dataset = dataset_ops.Dataset.from_tensors(0).repeat(None)
for _ in range(chain_length):
dataset = dataset.map(lambda x: x)
if optimize_dataset:
options = options_lib.Options()
options.experimental_optimization.apply_default_optimizations = False
options.experimental_optimization.map_fusion = 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.1",
"parameters": "%d.%s" % (chain_length, optimize_dataset),
},
name="map_fusion_{}_chain_length_{}".format(opt_mark, chain_length))
def benchmark_map_and_filter_fusion(self):
"""Evaluates performance map of fusion."""
chain_lengths = [0, 1, 2, 5, 10, 20, 50]
for chain_length in chain_lengths:
self._benchmark_map_and_filter_fusion(
chain_length=chain_length, optimize_dataset=False)
self._benchmark_map_and_filter_fusion(
chain_length=chain_length, optimize_dataset=True)
def _benchmark_map_and_filter_fusion(self, chain_length, optimize_dataset):
dataset = dataset_ops.Dataset.from_tensors(0).repeat(None)
for _ in range(chain_length):
dataset = dataset.map(lambda x: x + 5).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.map_and_filter_fusion = 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.2",
"parameters": "%d.%s" % (chain_length, optimize_dataset),
},
name="map_and_filter_fusion_{}_chain_length_{}".format(
opt_mark, chain_length))
# This benchmark compares the performance of pipeline with multiple chained
# filter with and without filter fusion.
def benchmark_filter_fusion(self):
chain_lengths = [0, 1, 2, 5, 10, 20, 50]
for chain_length in chain_lengths:
self._benchmark_filter_fusion(
chain_length=chain_length, optimize_dataset=False)
self._benchmark_filter_fusion(
chain_length=chain_length, optimize_dataset=True)
def _benchmark_filter_fusion(self, chain_length, optimize_dataset):
dataset = dataset_ops.Dataset.from_tensors(5).repeat(None)
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_fusion = 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.3",
"parameters": "%d.%s" % (chain_length, optimize_dataset),
},
name="filter_fusion_{}_chain_length_{}".format(opt_mark, chain_length))
# This benchmark compares the performance of pipeline with multiple chained
# filter with and without filter parallelization.
def benchmark_filter_parallelization(self):
chain_lengths = [0, 1, 2, 5, 10, 20, 50]
for chain_length in chain_lengths:
self._benchmark_filter_parallelization(
chain_length=chain_length, optimize_dataset=False)
self._benchmark_filter_parallelization(
chain_length=chain_length, optimize_dataset=True)
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))
if __name__ == "__main__":
benchmark_base.test.main()
@@ -0,0 +1,163 @@
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Benchmarks to compare effect of different parameter values on the performance."""
import time
import numpy as np
from tensorflow.python.data.benchmarks import benchmark_base
from tensorflow.python.data.experimental.ops import testing
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import script_ops
# The maximum sleeping time for each output step and the input sleeping time in
# milliseconds.
max_output_sleep_ms = 0.5
input_sleep_ms = 1.5
def sleep_function(x):
time.sleep(np.random.uniform(max_output_sleep_ms) / 1000)
return x
def map_function(x):
return script_ops.py_func(sleep_function, [x], x.dtype)
class ParameterValueBenchmark(benchmark_base.DatasetBenchmarkBase):
"""Benchmarks to compare effect of different parameter values on the performance."""
def _benchmark_map(self, num_parallel_calls, buffer_size):
k = 1024 * 1024
dataset = dataset_ops.Dataset.from_tensors(
(np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))).repeat()
dataset = dataset.map(
math_ops.matmul, num_parallel_calls=num_parallel_calls)
dataset = dataset.map(map_function)
dataset = dataset.prefetch(buffer_size=buffer_size)
dataset = dataset.apply(testing.sleep(int(input_sleep_ms * 1000)))
name_str = ("map_max_output_sleep_ms_%.2f_input_sleep_ms_%.2f_"
"num_parallel_calls_%d_buffer_size_%d")
return self.run_and_report_benchmark(
dataset=dataset,
num_elements=10000,
name=name_str %
(max_output_sleep_ms, input_sleep_ms, num_parallel_calls, buffer_size))
def benchmark_map(self):
nums_parallel_calls = [4, 8, 12]
buffer_sizes = [10, 50, 100, 150, 200, 250, 300]
parameters_list = []
wall_time_map = {}
for num_parallel_calls in nums_parallel_calls:
for buffer_size in buffer_sizes:
parameters = (num_parallel_calls, buffer_size)
parameters_list.append(parameters)
wall_time = self._benchmark_map(num_parallel_calls, buffer_size)
wall_time_map[parameters] = wall_time
parameters_list.sort(key=lambda x: wall_time_map[x])
for parameters in parameters_list:
print("num_parallel_calls_%d_buffer_size_%d_wall_time:" % parameters,
wall_time_map[parameters])
def _benchmark_map_and_batch(self, num_parallel_calls, buffer_size):
batch_size = 16
k = 1024 * 1024
dataset = dataset_ops.Dataset.from_tensors(
(np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))).repeat()
dataset = dataset.map(
math_ops.matmul, num_parallel_calls=num_parallel_calls)
dataset = dataset.batch(batch_size=batch_size)
dataset = dataset.map(map_function)
dataset = dataset.prefetch(buffer_size=buffer_size)
dataset = dataset.apply(testing.sleep(int(input_sleep_ms * 1000)))
name_str = ("map_and_batch_max_output_sleep_ms_%.2f_input_sleep_ms_%.2f"
"_num_parallel_calls_%d_buffer_size_%d")
return self.run_and_report_benchmark(
dataset=dataset,
num_elements=1000,
name=name_str %
(max_output_sleep_ms, input_sleep_ms, num_parallel_calls, buffer_size))
def benchmark_map_and_batch(self):
nums_parallel_calls = [4, 8, 12]
buffer_sizes = [10, 50, 100, 150, 200, 250, 300]
parameters_list = []
wall_time_map = {}
for num_parallel_calls in nums_parallel_calls:
for buffer_size in buffer_sizes:
parameters = (num_parallel_calls, buffer_size)
parameters_list.append(parameters)
wall_time = self._benchmark_map_and_batch(num_parallel_calls,
buffer_size)
wall_time_map[parameters] = wall_time
parameters_list.sort(key=lambda x: wall_time_map[x])
for parameters in parameters_list:
print("num_parallel_calls_%d_buffer_size_%d_wall_time:" % parameters,
wall_time_map[parameters])
def _benchmark_interleave(self, num_parallel_calls, buffer_size):
k = 1024 * 1024
dataset = dataset_ops.Dataset.from_tensors(
(np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))).repeat()
dataset = dataset.map(math_ops.matmul)
dataset = dataset.map(map_function)
dataset = dataset_ops.Dataset.range(1).repeat().interleave(
lambda _: dataset, # pylint: disable=cell-var-from-loop
cycle_length=10,
num_parallel_calls=num_parallel_calls)
dataset = dataset.prefetch(buffer_size=buffer_size)
dataset = dataset.apply(testing.sleep(int(input_sleep_ms * 1000)))
name_str = ("interleave_max_output_sleep_ms_%.2f_input_sleep_ms_%.2f"
"_num_parallel_calls_%d_buffer_size_%d")
return self.run_and_report_benchmark(
dataset=dataset,
num_elements=10000,
name=name_str %
(max_output_sleep_ms, input_sleep_ms, num_parallel_calls, buffer_size))
def benchmark_interleave(self):
nums_parallel_calls = [4, 8, 10]
buffer_sizes = [10, 50, 100, 150, 200, 250, 300]
parameters_list = []
wall_time_map = {}
for num_parallel_calls in nums_parallel_calls:
for buffer_size in buffer_sizes:
parameters = (num_parallel_calls, buffer_size)
parameters_list.append(parameters)
wall_time = self._benchmark_interleave(num_parallel_calls, buffer_size)
wall_time_map[parameters] = wall_time
parameters_list.sort(key=lambda x: wall_time_map[x])
for parameters in parameters_list:
print("num_parallel_calls_%d_buffer_size_%d_wall_time:" % parameters,
wall_time_map[parameters])
if __name__ == "__main__":
benchmark_base.test.main()
@@ -0,0 +1,65 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Benchmarks for `tf.data.experimental.rejection_resample()`."""
import numpy as np
from tensorflow.python.data.benchmarks import benchmark_base
from tensorflow.python.data.experimental.ops import resampling
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import options as options_lib
class RejectionResampleBenchmark(benchmark_base.DatasetBenchmarkBase):
"""Benchmarks for `tf.data.experimental.rejection_resample()`."""
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))
if __name__ == "__main__":
benchmark_base.test.main()
@@ -0,0 +1,200 @@
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Benchmarks for `tf.data.experimental.snapshot()`."""
import os
import shutil
from tensorflow.python.data.benchmarks import benchmark_base
from tensorflow.python.data.experimental.ops import snapshot
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.platform import test
class SnapshotDatasetBenchmark(benchmark_base.DatasetBenchmarkBase):
"""Benchmarks for `tf.data.experimental.snapshot()`."""
def _makeSnapshotDirectory(self):
tmp_dir = test.get_temp_dir()
tmp_dir = os.path.join(tmp_dir, "snapshot")
if os.path.exists(tmp_dir):
shutil.rmtree(tmp_dir)
os.mkdir(tmp_dir)
return tmp_dir
def _createSimpleDataset(self,
num_elements,
tmp_dir=None,
compression=snapshot.COMPRESSION_NONE):
if not tmp_dir:
tmp_dir = self._makeSnapshotDirectory()
dataset = dataset_ops.Dataset.from_tensor_slices([1.0])
dataset = dataset.map(
lambda x: gen_array_ops.broadcast_to(x, [50, 50, 3]))
dataset = dataset.repeat(num_elements)
dataset = dataset.apply(
snapshot.legacy_snapshot(tmp_dir, compression=compression))
return dataset
def benchmarkWriteSnapshotGzipCompression(self):
num_elements = 500000
dataset = self._createSimpleDataset(
num_elements=num_elements, compression=snapshot.COMPRESSION_GZIP)
self.run_and_report_benchmark(
dataset=dataset,
num_elements=num_elements,
name="write_gzip",
warmup=False,
extras={
"model_name": "snapshot.benchmark.1",
"parameters": "%d" % num_elements,
},
iters=1)
def benchmarkWriteSnapshotSnappyCompression(self):
num_elements = 500000
dataset = self._createSimpleDataset(
num_elements=num_elements, compression=snapshot.COMPRESSION_SNAPPY)
self.run_and_report_benchmark(
dataset=dataset,
num_elements=num_elements,
name="write_snappy",
warmup=False,
extras={
"model_name": "snapshot.benchmark.2",
"parameters": "%d" % num_elements,
},
iters=1)
def benchmarkWriteSnapshotSimple(self):
num_elements = 500000
dataset = self._createSimpleDataset(num_elements=num_elements)
# We only run one iteration here because running multiple iterations will
# cause the later iterations to simply read from the already written
# snapshot rather than write a new one.
self.run_and_report_benchmark(
dataset=dataset,
num_elements=num_elements,
name="write_simple",
warmup=False,
extras={
"model_name": "snapshot.benchmark.3",
"parameters": "%d" % num_elements,
},
iters=1)
def benchmarkPassthroughSnapshotSimple(self):
num_elements = 100000
tmp_dir = self._makeSnapshotDirectory()
dataset = self._createSimpleDataset(
num_elements=num_elements, tmp_dir=tmp_dir)
# Consume only 1 element, thus making sure we don't finalize.
self.run_benchmark(
dataset=dataset,
num_elements=1,
iters=1,
warmup=False,
apply_default_optimizations=True)
# Now run the actual benchmarks and report them
self.run_and_report_benchmark(
dataset=dataset,
num_elements=num_elements,
name="passthrough_simple",
extras={
"model_name": "snapshot.benchmark.4",
"parameters": "%d" % num_elements,
},
)
def benchmarkReadSnapshotSimple(self):
num_elements = 100000
tmp_dir = self._makeSnapshotDirectory()
dataset = self._createSimpleDataset(
num_elements=num_elements, tmp_dir=tmp_dir)
# consume all the elements to let snapshot write things to disk
self.run_benchmark(
dataset=dataset,
num_elements=num_elements,
iters=1,
warmup=False,
apply_default_optimizations=True)
# Now run the actual benchmarks and report them
self.run_and_report_benchmark(
dataset=dataset,
num_elements=num_elements,
name="read_simple",
extras={
"model_name": "snapshot.benchmark.5",
"parameters": "%d" % num_elements,
})
def benchmarkReadSnapshotGzipCompression(self):
num_elements = 100000
tmp_dir = self._makeSnapshotDirectory()
dataset = self._createSimpleDataset(
num_elements=num_elements,
tmp_dir=tmp_dir,
compression=snapshot.COMPRESSION_GZIP)
# consume all the elements to let snapshot write things to disk
self.run_benchmark(
dataset=dataset,
num_elements=num_elements,
iters=1,
warmup=False,
apply_default_optimizations=True)
# Now run the actual benchmarks and report them
self.run_and_report_benchmark(
dataset=dataset, num_elements=num_elements, name="read_gzip",
extras={
"model_name": "snapshot.benchmark.6",
"parameters": "%d" % num_elements,
})
def benchmarkReadSnapshotSnappyCompression(self):
num_elements = 100000
tmp_dir = self._makeSnapshotDirectory()
dataset = self._createSimpleDataset(
num_elements=num_elements,
tmp_dir=tmp_dir,
compression=snapshot.COMPRESSION_SNAPPY)
# consume all the elements to let snapshot write things to disk
self.run_benchmark(
dataset=dataset,
num_elements=num_elements,
iters=1,
warmup=False,
apply_default_optimizations=True)
# Now run the actual benchmarks and report them
self.run_and_report_benchmark(
dataset=dataset,
num_elements=num_elements,
name="read_snappy",
extras={
"model_name": "snapshot.benchmark.7",
"parameters": "%d" % num_elements,
})
if __name__ == "__main__":
benchmark_base.test.main()
@@ -0,0 +1,66 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Benchmarks for `tf.data.Dataset.unbatch()`."""
from tensorflow.python.data.benchmarks import benchmark_base
from tensorflow.python.data.ops import dataset_ops
class UnbatchBenchmark(benchmark_base.DatasetBenchmarkBase):
"""Benchmarks for `tf.data.Dataset.unbatch()`."""
def benchmark_native_unbatch(self):
batch_sizes = [1, 2, 5, 10, 20, 50]
num_elements = 10000
for batch_size in batch_sizes:
dataset = dataset_ops.Dataset.from_tensors("element").repeat(None)
dataset = dataset.batch(batch_size)
dataset = dataset.unbatch()
self.run_and_report_benchmark(
dataset=dataset,
num_elements=num_elements,
iters=5,
extras={
"model_name": "unbatch.benchmark.1",
"parameters": "%d" % batch_size,
},
name="native_batch_size_%d" % batch_size)
# Include a benchmark of the previous `unbatch()` implementation that uses
# a composition of more primitive ops. Eventually we'd hope to generate code
# that is as good in both cases.
def benchmark_old_unbatch_implementation(self):
batch_sizes = [1, 2, 5, 10, 20, 50]
num_elements = 10000
for batch_size in batch_sizes:
dataset = dataset_ops.Dataset.from_tensors("element").repeat(None)
dataset = dataset.batch(batch_size)
dataset = dataset.flat_map(dataset_ops.Dataset.from_tensor_slices)
self.run_and_report_benchmark(
dataset=dataset,
num_elements=num_elements,
iters=5,
extras={
"model_name": "unbatch.benchmark.2",
"parameters": "%d" % batch_size,
},
name="unfused_batch_size_%d" % batch_size)
if __name__ == "__main__":
benchmark_base.test.main()