wehub-resource-sync
149 lines
5.2 KiB
Python
149 lines
5.2 KiB
Python
# Copyright 2015 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 split and grad of split."""
|
|
|
|
import numpy as np
|
|
|
|
from tensorflow.core.protobuf import config_pb2
|
|
from tensorflow.python.client import session as session_lib
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.ops import array_ops
|
|
from tensorflow.python.ops import control_flow_ops
|
|
from tensorflow.python.ops import variables
|
|
from tensorflow.python.platform import benchmark
|
|
from tensorflow.python.platform import test
|
|
from tensorflow.python.platform import tf_logging as logging
|
|
|
|
|
|
def build_graph(device, input_shape, output_sizes, axis):
|
|
"""Build a graph containing a sequence of split operations.
|
|
|
|
Args:
|
|
device: string, the device to run on.
|
|
input_shape: shape of the input tensor.
|
|
output_sizes: size of each output along axis.
|
|
axis: axis to be split along.
|
|
|
|
Returns:
|
|
An array of tensors to run()
|
|
"""
|
|
with ops.device("/%s:0" % device):
|
|
# Use a variable to prevent constant folding.
|
|
inp = variables.Variable(array_ops.zeros(input_shape, dtype="float32"))
|
|
|
|
outputs = []
|
|
for _ in range(100):
|
|
outputs.extend(array_ops.split(inp, output_sizes, axis))
|
|
return outputs
|
|
|
|
|
|
class SplitBenchmark(test.Benchmark):
|
|
"""Benchmark split!"""
|
|
|
|
def _run_graph(self, device, output_shape, variable, num_outputs, axis):
|
|
"""Run the graph and print its execution time.
|
|
|
|
Args:
|
|
device: string, the device to run on.
|
|
output_shape: shape of each output tensors.
|
|
variable: whether or not the output shape should be fixed
|
|
num_outputs: the number of outputs to split the input into
|
|
axis: axis to be split
|
|
|
|
Returns:
|
|
The duration of the run in seconds.
|
|
"""
|
|
# Fallback to CPU if GPU is requested but not available.
|
|
if device == "gpu" and not test.is_gpu_available():
|
|
device = "cpu"
|
|
|
|
graph = ops.Graph()
|
|
with graph.as_default():
|
|
if not variable:
|
|
if axis == 0:
|
|
input_shape = [output_shape[0] * num_outputs, output_shape[1]]
|
|
sizes = [output_shape[0] for _ in range(num_outputs)]
|
|
else:
|
|
input_shape = [output_shape[0], output_shape[1] * num_outputs]
|
|
sizes = [output_shape[1] for _ in range(num_outputs)]
|
|
else:
|
|
sizes = np.random.randint(
|
|
low=max(1, output_shape[axis] - 2),
|
|
high=output_shape[axis] + 2,
|
|
size=num_outputs)
|
|
total_size = np.sum(sizes)
|
|
if axis == 0:
|
|
input_shape = [total_size, output_shape[1]]
|
|
else:
|
|
input_shape = [output_shape[0], total_size]
|
|
|
|
outputs = build_graph(device, input_shape, sizes, axis)
|
|
config = config_pb2.ConfigProto(graph_options=config_pb2.GraphOptions(
|
|
optimizer_options=config_pb2.OptimizerOptions(
|
|
opt_level=config_pb2.OptimizerOptions.L0)))
|
|
with session_lib.Session(graph=graph, config=config) as session:
|
|
logging.set_verbosity("info")
|
|
variables.global_variables_initializer().run()
|
|
bench = benchmark.TensorFlowBenchmark()
|
|
bench.run_op_benchmark(
|
|
session,
|
|
outputs,
|
|
mbs=input_shape[0] * input_shape[1] * 4 * 100 / 1e6,
|
|
extras={
|
|
"input_shape": input_shape,
|
|
"variable": variable,
|
|
"axis": axis,
|
|
},
|
|
)
|
|
|
|
def _get_name(self, device, input_shape, variable, num_outputs, axis):
|
|
return "split_%s_%s_%s_%d_%d" % (
|
|
device,
|
|
str(input_shape).replace(" ", ""),
|
|
str(variable),
|
|
num_outputs,
|
|
axis,
|
|
)
|
|
|
|
def benchmark_split(self):
|
|
print("Forward vs backward concat")
|
|
shapes = [[2000, 8], [8, 2000], [100, 18], [1000, 18], [10000, 18],
|
|
[100, 97], [1000, 97], [10000, 1], [1, 10000]]
|
|
axis_ = [1] # 0 is very fast because it doesn't actually do any copying
|
|
num_outputs = 100
|
|
variable = [False, True] # fixed input size or not
|
|
for shape in shapes:
|
|
for axis in axis_:
|
|
for v in variable:
|
|
self._run_graph("gpu", shape, v, num_outputs, axis)
|
|
|
|
def benchmark_many_splits(self):
|
|
"""Benchmark SplitV with large spatial dimensions and many outputs."""
|
|
h, w = 1234, 456
|
|
# Scenario 1: Many small splits.
|
|
# Total channels = 25.
|
|
self._run_graph("cpu", [h * w, 1], False, 25, axis=1)
|
|
|
|
def benchmark_few_large_splits(self):
|
|
"""Benchmark SplitV with large spatial dimensions and few large outputs."""
|
|
h, w = 1234, 456
|
|
# Scenario 2: Few large splits (same total elements)
|
|
# Total channels = 25, split into 5 pieces of 5 channels each.
|
|
self._run_graph("cpu", [h * w, 5], False, 5, axis=1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test.main()
|