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
+273
View File
@@ -0,0 +1,273 @@
load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "py_test")
load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
py_library(
name = "profiler",
srcs = ["profiler.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":model_analyzer",
":option_builder",
":tfprof_logger",
"//tensorflow/core/profiler:protos_all_py",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "profiler_client",
srcs = ["profiler_client.py"],
strict_deps = True,
deps = [
"//tensorflow/python/framework:errors",
"//tensorflow/python/profiler/internal:_pywrap_profiler_plugin",
"//tensorflow/python/util:tf_export",
],
)
cuda_py_strict_test(
name = "profiler_client_test",
srcs = ["profiler_client_test.py"],
tags = ["no_pip"],
deps = [
":profiler_client",
":profiler_v2",
"//tensorflow/python/eager:test",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:test_lib",
"@pypi//portpicker",
],
)
py_library(
name = "profiler_v2",
srcs = ["profiler_v2.py"],
strict_deps = True,
visibility = ["//tensorflow:internal"],
deps = [
"//tensorflow/python/framework:errors",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/profiler/internal:_pywrap_profiler",
"//tensorflow/python/util:tf_export",
],
)
cuda_py_strict_test(
name = "profiler_v2_test",
srcs = ["profiler_v2_test.py"],
tags = [
"no_pip",
],
deps = [
":profiler_v2",
":trace",
"//tensorflow/python/eager:test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:gfile",
],
)
py_test(
name = "profiler_wrapper_test",
srcs = ["profiler_wrapper_test.py"],
strict_deps = True,
tags = [
"no_pip",
],
deps = [
"//tensorflow/python/eager:test",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/profiler/internal:_pywrap_profiler_plugin",
],
)
py_library(
name = "option_builder",
srcs = ["option_builder.py"],
strict_deps = True,
deps = [
":tfprof_logger",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "model_analyzer",
srcs = ["model_analyzer.py"],
strict_deps = True,
deps = [
":option_builder",
":tfprof_logger",
"//tensorflow/core/profiler:protos_all_py",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/util:_pywrap_tfprof",
"//tensorflow/python/util:tf_export",
],
)
cuda_py_strict_test(
name = "model_analyzer_test",
srcs = ["model_analyzer_test.py"],
tags = [
"no_pip",
"notap",
"oss_serial",
],
xla_enable_strict_auto_jit = False, # Node names are different with autojit
deps = [
":model_analyzer",
":option_builder",
":profile_context",
"//tensorflow/core:protos_all_py",
"//tensorflow/core/profiler:protos_all_py",
"//tensorflow/python:distributed_framework_test_lib",
"//tensorflow/python/client:session",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:gradients",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/profiler/internal:model_analyzer_testlib",
"//tensorflow/python/util:compat",
],
)
cuda_py_strict_test(
name = "profiler_test",
srcs = ["profiler_test.py"],
tags = ["no_pip"],
xla_enable_strict_auto_jit = False, # Node names are different with autojit
deps = [
":model_analyzer",
":option_builder",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/profiler/internal:model_analyzer_testlib",
],
)
py_library(
name = "tfprof_logger",
srcs = ["tfprof_logger.py"],
strict_deps = True,
deps = [
"//tensorflow/core/profiler:protos_all_py",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/profiler/internal:flops_registry",
"//tensorflow/python/util:tf_export",
],
)
tf_py_strict_test(
name = "tfprof_logger_test",
size = "small",
srcs = ["tfprof_logger_test.py"],
deps = [
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
],
)
py_library(
name = "profile_context",
srcs = ["profile_context.py"],
strict_deps = True,
deps = [
":model_analyzer",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/util:_pywrap_tfprof",
"//tensorflow/python/util:compat",
],
)
cuda_py_strict_test(
name = "profile_context_test",
srcs = ["profile_context_test.py"],
tags = [
"no_gpu", # b/136036359
"no_pip",
],
xla_enable_strict_auto_jit = False, # Node names are different with autojit
deps = [
":option_builder",
":profile_context",
"//tensorflow/python/client:session",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/profiler/internal:model_analyzer_testlib",
],
)
py_library(
name = "pprof_profiler",
srcs = ["pprof_profiler.py"],
strict_deps = True,
deps = ["@com_google_pprof//:pprof_proto_py"],
)
py_test(
name = "pprof_profiler_test",
size = "small",
srcs = ["pprof_profiler_test.py"],
main = "pprof_profiler_test.py",
strict_deps = True,
tags = ["no_pip"], # TODO(annarev): get it working with pip.
deps = [
":pprof_profiler",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:client_testlib",
"@com_google_pprof//:pprof_proto_py",
],
)
py_library(
name = "trace",
srcs = ["trace.py"],
strict_deps = True,
visibility = [
"//perftools/accelerators/xprof/xprofilez/integration_tests:__pkg__",
"//tensorflow:internal",
],
deps = [
"//tensorflow/python/profiler/internal:_pywrap_traceme",
"//tensorflow/python/util:tf_export",
],
)
@@ -0,0 +1,39 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
py_library(
name = "mnist_testing_utils",
srcs = ["mnist_testing_utils.py"],
strict_deps = True,
deps = [
"//tensorflow:tensorflow_py",
"//tensorflow/python:extra_py_tests_deps",
],
)
cuda_py_strict_test(
name = "profiler_api_test",
srcs = ["profiler_api_test.py"],
tags = [
"no_oss", # TODO(b/283175845)
"no_pip",
"no_windows", # TODO(b/192257727)
],
deps = [
":mnist_testing_utils",
"//tensorflow/python/distribute:collective_all_reduce_strategy",
"//tensorflow/python/distribute:multi_process_runner",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/profiler:profiler_client",
"//tensorflow/python/profiler:profiler_v2",
"@pypi//portpicker",
],
)
@@ -0,0 +1,68 @@
# Copyright 2020 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.
# ==============================================================================
"""A simple MNIST model for testing multi-worker distribution strategies with Keras."""
import tensorflow as tf
def mnist_synthetic_dataset(batch_size, steps_per_epoch):
"""Generate synthetic MNIST dataset for testing."""
# train dataset
x_train = tf.ones([batch_size * steps_per_epoch, 28, 28, 1],
dtype=tf.dtypes.float32)
y_train = tf.ones([batch_size * steps_per_epoch, 1], dtype=tf.dtypes.int32)
train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_ds = train_ds.repeat()
# train_ds = train_ds.shuffle(100)
train_ds = train_ds.batch(64, drop_remainder=True)
# eval dataset
x_test = tf.random.uniform([10000, 28, 28, 1], dtype=tf.dtypes.float32)
y_test = tf.random.uniform([10000, 1],
minval=0,
maxval=9,
dtype=tf.dtypes.int32)
eval_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))
eval_ds = eval_ds.batch(64, drop_remainder=True)
return train_ds, eval_ds
def get_mnist_model(input_shape):
"""Define a deterministically-initialized CNN model for MNIST testing."""
inputs = tf.keras.Input(shape=input_shape)
x = tf.keras.layers.Conv2D(
32,
kernel_size=(3, 3),
activation="relu",
kernel_initializer=tf.keras.initializers.TruncatedNormal(seed=99))(
inputs)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.Flatten()(x) + tf.keras.layers.Flatten()(x)
x = tf.keras.layers.Dense(
10,
activation="softmax",
kernel_initializer=tf.keras.initializers.TruncatedNormal(seed=99))(
x)
model = tf.keras.Model(inputs=inputs, outputs=x)
# TODO(yuefengz): optimizer with slot variables doesn't work because of
# optimizer's bug.
# TODO(yuefengz): we should not allow non-v2 optimizer.
model.compile(
loss=tf.keras.losses.sparse_categorical_crossentropy,
optimizer=tf.keras.optimizers.SGD(learning_rate=0.001),
metrics=["accuracy"])
return model
@@ -0,0 +1,148 @@
# Copyright 2020 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.
# ==============================================================================
"""Tests for tf 2.x profiler."""
import glob
import os
import threading
import portpicker
from tensorflow.python.distribute import collective_all_reduce_strategy as collective_strategy
from tensorflow.python.distribute import multi_process_runner
from tensorflow.python.eager import context
from tensorflow.python.framework import test_util
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.profiler import profiler_client
from tensorflow.python.profiler import profiler_v2 as profiler
from tensorflow.python.profiler.integration_test import mnist_testing_utils
def _model_setup():
"""Set up a MNIST Keras model for testing purposes.
Builds a MNIST Keras model and returns model information.
Returns:
A tuple of (batch_size, steps, train_dataset, mode)
"""
context.set_log_device_placement(True)
batch_size = 64
steps = 2
with collective_strategy.CollectiveAllReduceStrategy().scope():
# TODO(b/142509827): In rare cases this errors out at C++ level with the
# "Connect failed" error message.
train_ds, _ = mnist_testing_utils.mnist_synthetic_dataset(batch_size, steps)
model = mnist_testing_utils.get_mnist_model((28, 28, 1))
return batch_size, steps, train_ds, model
def _make_temp_log_dir(test_obj):
return test_obj.get_temp_dir()
class ProfilerApiTest(test_util.TensorFlowTestCase):
def setUp(self):
super().setUp()
self.worker_start = threading.Event()
self.profile_done = False
def _check_xspace_pb_exist(self, logdir):
path = os.path.join(logdir, 'plugins', 'profile', '*', '*.xplane.pb')
self.assertEqual(1, len(glob.glob(path)),
'Expected one path match: ' + path)
def test_single_worker_no_profiling(self):
"""Test single worker without profiling."""
_, steps, train_ds, model = _model_setup()
model.fit(x=train_ds, epochs=2, steps_per_epoch=steps)
def test_single_worker_sampling_mode(self, delay_ms=None):
"""Test single worker sampling mode."""
def on_worker(port, worker_start):
logging.info('worker starting server on {}'.format(port))
profiler.start_server(port)
_, steps, train_ds, model = _model_setup()
worker_start.set()
while True:
model.fit(x=train_ds, epochs=2, steps_per_epoch=steps)
if self.profile_done:
break
def on_profile(port, logdir, worker_start):
# Request for 30 milliseconds of profile.
duration_ms = 30
worker_start.wait()
options = profiler.ProfilerOptions(
host_tracer_level=2,
python_tracer_level=0,
device_tracer_level=1,
delay_ms=delay_ms,
)
profiler_client.trace('localhost:{}'.format(port), logdir, duration_ms,
'', 100, options)
self.profile_done = True
logdir = self.get_temp_dir()
port = portpicker.pick_unused_port()
thread_profiler = threading.Thread(
target=on_profile, args=(port, logdir, self.worker_start))
thread_worker = threading.Thread(
target=on_worker, args=(port, self.worker_start))
thread_worker.start()
thread_profiler.start()
thread_profiler.join()
thread_worker.join(120)
self._check_xspace_pb_exist(logdir)
def test_single_worker_sampling_mode_short_delay(self):
"""Test single worker sampling mode with a short delay.
Expect that requested delayed start time will arrive late, and a subsequent
retry will issue an immediate start.
"""
self.test_single_worker_sampling_mode(delay_ms=1)
def test_single_worker_sampling_mode_long_delay(self):
"""Test single worker sampling mode with a long delay."""
self.test_single_worker_sampling_mode(delay_ms=1000)
def test_single_worker_programmatic_mode(self):
"""Test single worker programmatic mode."""
logdir = self.get_temp_dir()
options = profiler.ProfilerOptions(
host_tracer_level=2,
python_tracer_level=0,
device_tracer_level=1,
)
profiler.start(logdir, options)
_, steps, train_ds, model = _model_setup()
model.fit(x=train_ds, epochs=2, steps_per_epoch=steps)
profiler.stop()
self._check_xspace_pb_exist(logdir)
if __name__ == '__main__':
multi_process_runner.test_main()
+263
View File
@@ -0,0 +1,263 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("@xla//xla/tsl:tsl.bzl", "if_macos", "internal_visibility")
load("@xla//xla/tsl:tsl.default.bzl", "tsl_pybind_extension")
load("//tensorflow:tensorflow.bzl", "py_test")
load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "get_compatible_with_portable", "tf_python_pybind_extension")
load("//tensorflow/core/profiler/builds:build_config.bzl", "tf_profiler_copts")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//tensorflow/python/profiler:__subpackages__"],
licenses = ["notice"],
)
py_library(
name = "flops_registry",
srcs = ["flops_registry.py"],
strict_deps = True,
deps = [
"//tensorflow/python/framework:graph_util",
"//tensorflow/python/framework:ops",
"//third_party/py/numpy",
],
)
py_test(
name = "flops_registry_test",
srcs = ["flops_registry_test.py"],
strict_deps = True,
deps = [
":flops_registry",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
],
)
py_library(
name = "model_analyzer_testlib",
srcs = ["model_analyzer_testlib.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:init_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn_grad",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:rnn",
"//tensorflow/python/ops:rnn_cell",
"//tensorflow/python/ops:tensor_array_grad",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/profiler:model_analyzer",
"//tensorflow/python/training:gradient_descent",
"//tensorflow/python/util:_pywrap_tfprof",
"//tensorflow/python/util:compat",
],
)
py_test(
name = "print_model_analysis_test",
srcs = ["print_model_analysis_test.py"],
strict_deps = True,
deps = [
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:init_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "run_metadata_test",
srcs = ["run_metadata_test.py"],
tags = [
"no_gpu", # b/138442728
"no_pip",
],
xla_enable_strict_auto_jit = False, # Node names are different with autojit
deps = [
":model_analyzer_testlib",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/profiler:model_analyzer",
"//tensorflow/python/profiler:option_builder",
],
)
tf_python_pybind_extension(
name = "_pywrap_traceme",
srcs = ["traceme_wrapper.cc"],
copts = tf_profiler_copts(),
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_traceme.pyi",
],
visibility = [
"//perftools/accelerators/xprof/xprofilez/integration_tests:__pkg__",
"//tensorflow/python:__pkg__",
"//tensorflow/python/profiler:__subpackages__",
"//tensorflow/tools/pip_package:__subpackages__",
],
deps = [
"@pybind11",
"@xla//xla/python/profiler/internal:traceme_state",
"@xla//xla/python/profiler/internal:traceme_wrapper",
],
)
tf_python_pybind_extension(
name = "_pywrap_profiler",
srcs = ["profiler_wrapper.cc"],
copts = tf_profiler_copts(),
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_profiler.pyi",
],
visibility = [
"//tensorflow/core/profiler:internal",
"//tensorflow/python:__pkg__",
"//tensorflow/python/eager:__pkg__",
"//tensorflow/python/profiler:__pkg__",
"//tensorflow/tools/pip_package:__subpackages__",
],
deps = [
":profiler_pywrap_impl",
"//tensorflow/core/profiler/rpc:profiler_server_for_pybind",
"//tensorflow/python/lib/core:pybind11_status",
"@com_google_absl//absl/status",
"@org_xprof//xprof/convert:repository",
"@org_xprof//xprof/convert:tool_options",
"@org_xprof//xprof/convert:xplane_to_tools_data",
"@pybind11",
],
)
cc_library(
name = "python_hooks",
hdrs = ["python_hooks.h"],
compatible_with = get_compatible_with_portable(),
copts = tf_profiler_copts() + ["-fexceptions"],
features = ["-use_header_modules"], # Incompatible with -fexceptions.
visibility = ["//visibility:private"],
deps = [
"//tensorflow/core:lib",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@pybind11",
"@tsl//tsl/profiler/protobuf:xplane_proto_cc",
"@xla//xla/python/profiler/internal:python_hooks",
],
alwayslink = True,
)
cc_library(
name = "profiler_pywrap_impl",
srcs = ["profiler_pywrap_impl.cc"],
hdrs = ["profiler_pywrap_impl.h"],
visibility = ["//visibility:private"],
deps = [
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core/profiler/lib:profiler_session_for_pybind",
"//tensorflow/core/profiler/rpc:profiler_server_for_pybind",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:variant",
"@tsl//tsl/profiler/lib:profiler_session",
"@tsl//tsl/profiler/protobuf:xplane_proto_cc",
"@xla//xla/tsl/profiler/convert:xplane_to_trace_events",
"@xla//xla/tsl/profiler/rpc/client:capture_profile",
"@xla//xla/tsl/profiler/utils:session_manager",
],
)
tsl_pybind_extension(
name = "_pywrap_profiler_plugin",
srcs = ["pywrap_profiler_plugin.cc"],
pytype_srcs = [
"_pywrap_profiler_plugin.pyi",
],
visibility = internal_visibility([
"//tensorflow/core/profiler:internal",
"//tensorflow/python/eager:__pkg__",
"//tensorflow/python/profiler:__pkg__",
]),
deps = [
"//tensorflow/core:protos_all_cc_impl",
"//tensorflow/core/framework:attr_value_proto_cc_impl",
"//tensorflow/core/framework:op",
"//tensorflow/core/framework:tensor",
"//tensorflow/python/lib/core:py_exception_registry",
"//tensorflow/python/lib/core:pybind11_status",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_protobuf//:protobuf",
"@org_xprof//xprof/convert:tool_options",
"@org_xprof//xprof/pywrap:profiler_plugin_impl",
"@pybind11",
"@tsl//tsl/profiler/protobuf:profiler_analysis_proto_cc_impl",
"@tsl//tsl/profiler/protobuf:profiler_options_proto_cc_impl",
"@tsl//tsl/profiler/protobuf:profiler_service_monitor_result_proto_cc_impl",
"@tsl//tsl/profiler/protobuf:profiler_service_proto_cc_impl",
"@xla//xla:autotune_results_proto_cc_impl",
"@xla//xla:autotuning_proto_cc_impl",
"@xla//xla:xla_data_proto_cc_impl",
"@xla//xla:xla_proto_cc_impl",
"@xla//xla/pjrt:status_casters",
"@xla//xla/service:hlo_proto_cc_impl",
"@xla//xla/service:metrics_proto_cc_impl",
"@xla//xla/service/gpu:backend_configs_cc_impl",
"@xla//xla/service/gpu/model:hlo_op_profile_proto_cc_impl",
"@xla//xla/stream_executor:device_description_proto_cc_impl",
"@xla//xla/stream_executor/cuda:cuda_compute_capability_proto_cc_impl",
"@xla//xla/tsl/framework:allocator_registry_impl",
"@xla//xla/tsl/lib/io:table",
"@xla//xla/tsl/platform:env_impl",
"@xla//xla/tsl/platform:types",
"@xla//xla/tsl/platform/cloud:gcs_file_system",
"@xla//xla/tsl/profiler/backends/cpu:traceme_recorder_impl",
"@xla//xla/tsl/profiler/rpc:profiler_server_impl",
"@xla//xla/tsl/profiler/rpc:profiler_service_impl",
"@xla//xla/tsl/profiler/rpc/client:capture_profile",
"@xla//xla/tsl/profiler/rpc/client:profiler_client_impl",
"@xla//xla/tsl/profiler/utils:session_manager",
"@xla//xla/tsl/profiler/utils:time_utils_impl",
"@xla//xla/tsl/protobuf:bfc_memory_map_proto_cc_impl",
"@xla//xla/tsl/protobuf:dnn_proto_cc_impl",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc_impl",
"@xla//xla/tsl/protobuf:histogram_proto_cc_impl",
"@xla//xla/tsl/protobuf:rpc_options_proto_cc_impl",
"@xla//xla/tsl/protobuf:test_log_proto_cc_impl",
] + if_macos([
"@xla//xla/tsl/lib/histogram",
"@xla//xla/tsl/lib/io:record_writer",
"@xla//xla/tsl/lib/monitoring:collection_registry",
"@xla//xla/tsl/lib/monitoring:sampler",
"//tensorflow/core:framework_internal_impl",
"//tensorflow/core/common_runtime/gpu:gpu_id_impl",
"//tensorflow/core/lib/core:arena",
"//tensorflow/core/lib/strings:ordered_code",
"//tensorflow/core/platform:platform_strings",
"@tsl//tsl/platform:random",
"@xla//xla/tsl/platform:resource",
"@tsl//tsl/profiler/lib:profiler_factory_impl",
"@tsl//tsl/profiler/lib:profiler_session_impl",
]),
)
@@ -0,0 +1,22 @@
# Copyright 2023 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.
# ==============================================================================
class ProfilerSession:
def __init__(self) -> None: ...
def export_to_tb(self) -> None: ...
def start(self, arg0: str, arg1: dict) -> None: ...
def stop(self) -> bytes: ...
def start_server(arg0: int) -> None: ...
@@ -0,0 +1,19 @@
# Copyright 2023 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.
# ==============================================================================
def monitor(arg0: str, arg1: int, arg2: int, arg3: bool) -> str: ...
def trace(arg0: str, arg1: str, arg2: str, arg3: bool, arg4: int, arg5: int, arg6: dict) -> None: ...
def xspace_to_tools_data(arg0: list, arg1: str, arg2: dict = ...) -> tuple: ...
def xspace_to_tools_data_from_byte_string(arg0: list, arg1: list, arg2: str, arg3: dict) -> tuple: ...
@@ -0,0 +1,21 @@
# Copyright 2023 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.
# ==============================================================================
class TraceMe:
def __init__(self, arg0: str, **kwargs) -> None: ...
def SetMetadata(self, **kwargs) -> None: ...
def Stop(self) -> None: ...
def traceme_enabled(*args, **kwargs): ...
@@ -0,0 +1,480 @@
# Copyright 2016 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.
# ==============================================================================
"""Register flops statistics for various TensorFlow operations.
"""
import numpy as np
from tensorflow.python.framework import graph_util
from tensorflow.python.framework import ops
# List of all ops which have implemented flops statistics.
IMPLEMENTED_OPS = set([
# Unary ops
"Reciprocal", "Square", "Rsqrt", "Log", "Neg", "AssignSub", "AssignAdd",
"L2Loss", "Softmax",
# Binary ops
"Add", "Sub", "Mul", "RealDiv", "Maximum", "Minimum", "Pow", "RsqrtGrad",
"GreaterEqual", "Greater", "LessEqual", "Less", "Equal", "NotEqual",
"SquaredDifference", "AddV2",
# Reduction ops
"Mean", "Sum", "ArgMax", "ArgMin", "BiasAddGrad",
# Convolution and pooling
"AvgPool", "MaxPool", "AvgPoolGrad", "MaxPoolGrad", "Conv2DBackpropInput",
"Conv2DBackpropFilter",
# Other ops
"AddN", "MatMul",
# Ops implemented in core tensorflow:
"Conv2D", "DepthwiseConv2dNative", "BiasAdd", "Dilation2D",
])
def _zero_flops(graph, node):
"""Returns zero flops."""
del graph, node # graph and node are unused
return ops.OpStats("flops", 0)
def _list_product(lst):
"""Computes product of element of the list."""
result = 1
for item in lst:
result *= item
return result
################################################################################
# Unary operations
################################################################################
def _unary_op_flops(graph, node, ops_per_element=1):
"""Common code which compute flops for unary operations."""
in_shape = graph_util.tensor_shape_from_node_def_name(graph, node.input[0])
in_shape.assert_is_fully_defined()
return ops.OpStats("flops", in_shape.num_elements() * ops_per_element)
@ops.RegisterStatistics("Reciprocal", "flops")
def _reciprocal_flops(graph, node):
"""Compute flops for Reciprocal operation."""
return _unary_op_flops(graph, node)
@ops.RegisterStatistics("Square", "flops")
def _square_flops(graph, node):
"""Compute flops for Square operation."""
return _unary_op_flops(graph, node)
@ops.RegisterStatistics("Rsqrt", "flops")
def _rsqrt_flops(graph, node):
"""Compute flops for Rsqrt operation."""
# Rsqrt(x) = 1 / sqrt(x)
return _unary_op_flops(graph, node, ops_per_element=2)
@ops.RegisterStatistics("Log", "flops")
def _log_flops(graph, node):
"""Compute flops for Log operation."""
return _unary_op_flops(graph, node)
@ops.RegisterStatistics("Neg", "flops")
def _neg_flops(graph, node):
"""Compute flops for Neg operation."""
return _unary_op_flops(graph, node)
@ops.RegisterStatistics("AssignSub", "flops")
def _assign_sub_flops(graph, node):
"""Compute flops for AssignSub operation."""
return _unary_op_flops(graph, node)
@ops.RegisterStatistics("AssignAdd", "flops")
def _assign_add_flops(graph, node):
"""Compute flops for AssignAdd operation."""
return _unary_op_flops(graph, node)
@ops.RegisterStatistics("L2Loss", "flops")
def _l2_loss_flops(graph, node):
"""Compute flops for L2Loss operation."""
in_shape = graph_util.tensor_shape_from_node_def_name(graph, node.input[0])
in_shape.assert_is_fully_defined()
# Tensorflow uses inefficient implementation, with (3*N-1) flops:
# Optimal implementation is 2*N flops
return ops.OpStats("flops", in_shape.num_elements() * 3 - 1)
@ops.RegisterStatistics("Softmax", "flops")
def _softmax_flops(graph, node):
"""Compute flops for Softmax operation."""
# Softmax implementation:
#
# Approximate flops breakdown:
# 2*n -- compute shifted logits
# n -- exp of shifted logits
# 2*n -- compute softmax from exp of shifted logits
return _unary_op_flops(graph, node, ops_per_element=5)
################################################################################
# Binary operations
################################################################################
def _binary_per_element_op_flops(graph, node, ops_per_element=1):
"""Common code which compute flops for binary operations."""
out_shape = graph_util.tensor_shape_from_node_def_name(graph, node.name)
out_shape.assert_is_fully_defined()
return ops.OpStats("flops", out_shape.num_elements() * ops_per_element)
@ops.RegisterStatistics("Add", "flops")
@ops.RegisterStatistics("AddV2", "flops")
def _add_flops(graph, node):
"""Compute flops for Add operation."""
return _binary_per_element_op_flops(graph, node)
@ops.RegisterStatistics("Sub", "flops")
def _sub_flops(graph, node):
"""Compute flops for Sub operation."""
return _binary_per_element_op_flops(graph, node)
@ops.RegisterStatistics("Mul", "flops")
def _mul_flops(graph, node):
"""Compute flops for Mul operation."""
return _binary_per_element_op_flops(graph, node)
@ops.RegisterStatistics("RealDiv", "flops")
def _real_div_flops(graph, node):
"""Compute flops for RealDiv operation."""
return _binary_per_element_op_flops(graph, node)
@ops.RegisterStatistics("Maximum", "flops")
def _maximum_flops(graph, node):
"""Compute flops for Maximum operation."""
return _binary_per_element_op_flops(graph, node)
@ops.RegisterStatistics("Minimum", "flops")
def _minimum_flops(graph, node):
"""Compute flops for Minimum operation."""
return _binary_per_element_op_flops(graph, node)
@ops.RegisterStatistics("Pow", "flops")
def _pow_flops(graph, node):
"""Compute flops for Pow operation."""
return _binary_per_element_op_flops(graph, node)
@ops.RegisterStatistics("RsqrtGrad", "flops")
def _rsqrt_grad_flops(graph, node):
"""Compute flops for RsqrtGrad operation."""
return _binary_per_element_op_flops(graph, node, ops_per_element=4)
@ops.RegisterStatistics("GreaterEqual", "flops")
def _greater_equal_flops(graph, node):
"""Compute flops for GreaterEqual operation."""
return _binary_per_element_op_flops(graph, node)
@ops.RegisterStatistics("Greater", "flops")
def _greater_flops(graph, node):
"""Compute flops for Greater operation."""
return _binary_per_element_op_flops(graph, node)
@ops.RegisterStatistics("LessEqual", "flops")
def _less_equal_flops(graph, node):
"""Compute flops for LessEqual operation."""
return _binary_per_element_op_flops(graph, node)
@ops.RegisterStatistics("Less", "flops")
def _less_flops(graph, node):
"""Compute flops for Less operation."""
return _binary_per_element_op_flops(graph, node)
@ops.RegisterStatistics("Equal", "flops")
def _equal_flops(graph, node):
"""Compute flops for Equal operation."""
return _binary_per_element_op_flops(graph, node)
@ops.RegisterStatistics("NotEqual", "flops")
def _not_equal_flops(graph, node):
"""Compute flops for NotEqual operation."""
return _binary_per_element_op_flops(graph, node)
@ops.RegisterStatistics("SquaredDifference", "flops")
def _squared_difference_flops(graph, node):
"""Compute flops for SquaredDifference operation."""
return _binary_per_element_op_flops(graph, node, ops_per_element=2)
################################################################################
# Reduction ops
################################################################################
def _reduction_op_flops(graph, node, reduce_flops=1, finalize_flops=0):
"""Common code which compute flops for reduction operations."""
in_shape = graph_util.tensor_shape_from_node_def_name(graph, node.input[0])
in_shape.assert_is_fully_defined()
out_shape = graph_util.tensor_shape_from_node_def_name(graph, node.name)
out_shape.assert_is_fully_defined()
num_flops = (in_shape.num_elements() * reduce_flops
+ out_shape.num_elements() * (finalize_flops - reduce_flops))
return ops.OpStats("flops", num_flops)
@ops.RegisterStatistics("Mean", "flops")
def _mean_flops(graph, node):
"""Compute flops for Mean operation."""
# reduction - sum, finalization - divide
return _reduction_op_flops(graph, node, reduce_flops=1, finalize_flops=1)
@ops.RegisterStatistics("Sum", "flops")
def _sum_flops(graph, node):
"""Compute flops for Sum operation."""
# reduction - sum, no finalization
return _reduction_op_flops(graph, node, reduce_flops=1, finalize_flops=0)
@ops.RegisterStatistics("ArgMax", "flops")
def _arg_max_flops(graph, node):
"""Compute flops for ArgMax operation."""
# reduction - comparison, no finalization
return _reduction_op_flops(graph, node, reduce_flops=1, finalize_flops=0)
@ops.RegisterStatistics("ArgMin", "flops")
def _arg_min_flops(graph, node):
"""Compute flops for ArgMin operation."""
# reduction - comparison, no finalization
return _reduction_op_flops(graph, node, reduce_flops=1, finalize_flops=0)
@ops.RegisterStatistics("BiasAddGrad", "flops")
def _bias_add_grad_flops(graph, node):
"""Compute flops for BiasAddGrad operation."""
# Implementation of BiasAddGrad, essentially it's a reduce sum and reshaping:
# So computing flops same way as for "Sum"
return _reduction_op_flops(graph, node, reduce_flops=1, finalize_flops=0)
################################################################################
# Convolution and pooling
# Note: all flops statistics are implemented only for NHWC data format
################################################################################
def _verify_conv_data_format(node):
"""Verifies data format for pooling and convolutional operations."""
# TODO(xpan): P1: Support NCHW
if node.attr["data_format"].s != b"NHWC":
raise ValueError("Only NHWC format is supported in flops computations")
def _pool_flops(graph, node):
"""Common code which compute flops for pooling operations."""
# compute flops for average and max pooling
_verify_conv_data_format(node)
#
# Pooling declaration:
# Inputs:
# - value
# Outputs:
# - output
# Attributes:
# - ksize
# - strides
# - padding
# - data_format
#
# Pooling implementation:
out_shape = graph_util.tensor_shape_from_node_def_name(graph, node.name)
out_shape.assert_is_fully_defined()
kernel_shape = list(node.attr["ksize"].list.i)
kernel_area = _list_product(kernel_shape)
return ops.OpStats("flops", kernel_area * out_shape.num_elements())
@ops.RegisterStatistics("AvgPool", "flops")
def _avg_pool_flops(graph, node):
"""Compute flops for AvgPool operation."""
return _pool_flops(graph, node)
@ops.RegisterStatistics("MaxPool", "flops")
def _max_pool_flops(graph, node):
"""Compute flops for MaxPool operation."""
return _pool_flops(graph, node)
@ops.RegisterStatistics("AvgPoolGrad", "flops")
def _avg_pool_grad_flops(graph, node):
"""Compute flops for AvgPoolGrad operation."""
_verify_conv_data_format(node)
# Pooling gradient implementation:
out_backprop_shape = graph_util.tensor_shape_from_node_def_name(graph,
node.input[1])
out_backprop_shape.assert_is_fully_defined()
kernel_shape = list(node.attr["ksize"].list.i)
kernel_area = _list_product(kernel_shape)
# TensorFlow multiply each element of pooling window by coefficient,
# then sum up all of them, thus we have 2 flops per element:
# More optimal implementation - if division is done after.
return ops.OpStats("flops",
kernel_area * out_backprop_shape.num_elements() * 2)
@ops.RegisterStatistics("MaxPoolGrad", "flops")
def _max_pool_grad_flops(graph, node):
"""Compute flops for MaxPoolGrad operation."""
_verify_conv_data_format(node)
#
# MaxPoolGrad declaration:
# Inputs:
# - orig_input -- original input tensor (of max_pool)
# - orig_output -- original output tensor (of max_pool)
# - grad -- gradient with respect to output of max_pool
# Outputs:
# - output -- gradient with respect to input of max_pool
# Attributes:
# - ksize
# - strides
# - padding
# - data_format
# It computes MaxPool first, then one flop per each element of original output
#
kernel_shape = list(node.attr["ksize"].list.i)
kernel_area = _list_product(kernel_shape)
orig_out_shape = graph_util.tensor_shape_from_node_def_name(graph,
node.input[1])
orig_out_shape.assert_is_fully_defined()
max_pool_ops = kernel_area * orig_out_shape.num_elements()
return ops.OpStats("flops", max_pool_ops + orig_out_shape.num_elements())
@ops.RegisterStatistics("Conv2DBackpropInput", "flops")
def _conv_2d_backprop_input_flops(graph, node):
"""Compute flops for Conv2DBackpropInput operation."""
# Formula:
# batch_size * image_x_dim * image_y_dim * kernel_x_dim * kernel_y_dim
# * input_depth * output_depth * 2 / (image_x_stride * image_x_stride)
#
# Where:
# image_x_dim, image_y_dim and input_depth --- size of input to source (no
# backprop) convolution, in other words they are sizes of backprop output.
# output_depth --- number of filters in the original convolution, thus
# depth of backprop input.
# kernel_x_dim and kernel_y_dim --- sizes of filter in spatial dimension
# image_x_stride and image_x_stride --- strides of the convolution
#
_verify_conv_data_format(node)
# out_shape = [batch_size, image_y_dim, image_x_dim, input_depth]
out_shape = graph_util.tensor_shape_from_node_def_name(graph, node.name)
out_shape.assert_is_fully_defined()
# kernel_shape = [kernel_y_dim, kernel_x_dim, input_depth, output_depth]
kernel_shape = graph_util.tensor_shape_from_node_def_name(graph,
node.input[1])
kernel_shape.assert_is_fully_defined()
# strides
strides_shape = list(node.attr["strides"].list.i)
strides_product = strides_shape[1] * strides_shape[2]
return ops.OpStats("flops",
(2 * out_shape.num_elements()
* kernel_shape.num_elements()
/ (out_shape.dims[-1].value * strides_product)))
@ops.RegisterStatistics("Conv2DBackpropFilter", "flops")
def _conv_2d_backprop_filter_flops(graph, node):
"""Compute flops for Conv2DBackpropFilter operation."""
# Formula same as for Conv2DBackpropInput:
# batch_size * image_x_dim * image_y_dim * kernel_x_dim * kernel_y_dim
# * input_depth * output_depth * 2 / (image_x_stride * image_x_stride)
#
_verify_conv_data_format(node)
# image_shape = [batch_size, image_y_dim, image_x_dim, input_depth]
image_shape = graph_util.tensor_shape_from_node_def_name(graph, node.input[0])
image_shape.assert_is_fully_defined()
# kernel_shape = [kernel_y_dim, kernel_x_dim, input_depth, output_depth]
kernel_shape = graph_util.tensor_shape_from_node_def_name(graph, node.name)
kernel_shape.assert_is_fully_defined()
# strides
strides_shape = list(node.attr["strides"].list.i)
strides_product = strides_shape[1] * strides_shape[2]
return ops.OpStats("flops",
(2 * image_shape.num_elements()
* kernel_shape.num_elements()
/ (image_shape.dims[-1].value * strides_product)))
################################################################################
# Other ops
################################################################################
@ops.RegisterStatistics("AddN", "flops")
def _add_n_flops(graph, node):
"""Compute flops for AddN operation."""
if not node.input:
return _zero_flops(graph, node)
in_shape = graph_util.tensor_shape_from_node_def_name(graph, node.input[0])
in_shape.assert_is_fully_defined()
return ops.OpStats("flops", in_shape.num_elements() * (len(node.input) - 1))
@ops.RegisterStatistics("MatMul", "flops")
def _calc_mat_mul_flops(graph, node):
"""Calculates the compute resources needed for MatMul."""
transpose_a = node.attr["transpose_a"].b
a_shape = graph_util.tensor_shape_from_node_def_name(graph, node.input[0])
a_shape.assert_is_fully_defined()
if transpose_a:
k = int(a_shape[0])
else:
k = int(a_shape[1])
output_shape = graph_util.tensor_shape_from_node_def_name(graph, node.name)
output_shape.assert_is_fully_defined()
output_count = np.prod(output_shape.as_list())
return ops.OpStats("flops", (k * output_count * 2))
@ops.RegisterStatistics("BatchMatMul", "flops")
@ops.RegisterStatistics("BatchMatMulV2", "flops")
@ops.RegisterStatistics("BatchMatMulV3", "flops")
def _calc_batch_mat_mul_flops(graph, node):
"""Calculates the compute resources needed for BatchMatMul."""
transpose_a = node.attr["transpose_a"].b
a_shape = graph_util.tensor_shape_from_node_def_name(graph, node.input[0])
a_shape.assert_is_fully_defined()
if transpose_a:
k = int(a_shape[-2])
else:
k = int(a_shape[-1])
output_shape = graph_util.tensor_shape_from_node_def_name(graph, node.name)
output_shape.assert_is_fully_defined()
output_count = np.prod(output_shape.as_list())
return ops.OpStats("flops", (k * output_count * 2))
@@ -0,0 +1,51 @@
# Copyright 2023 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.
# ==============================================================================
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.profiler.internal import flops_registry # pylint: disable=unused-import
class FlopsRegistryTest(test.TestCase):
@test_util.run_v1_only('Test requires a Graph and NodeDef inspection')
def testSimpleStatistics(self):
a = variables.Variable(random_ops.random_normal([25, 16]))
b = variables.Variable(random_ops.random_normal([16, 9]))
math_ops.matmul(a, b)
g = ops.get_default_graph()
for op in g.get_operations():
flops = ops.get_stats_for_node_def(g, op.node_def, 'flops').value
if op.name == 'MatMul':
self.assertEqual(7200, flops)
@test_util.run_v1_only('Test requires a Graph and NodeDef inspection')
def testTransposedStatistics(self):
a = variables.Variable(random_ops.random_normal([16, 25]))
b = variables.Variable(random_ops.random_normal([16, 9]))
math_ops.matmul(a, b, transpose_a=True)
g = ops.get_default_graph()
for op in g.get_operations():
flops = ops.get_stats_for_node_def(g, op.node_def, 'flops').value
if op.name == 'MatMul':
self.assertEqual(7200, flops)
if __name__ == '__main__':
test.main()
@@ -0,0 +1,113 @@
# Copyright 2016 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.
# ==============================================================================
"""A test lib that defines some models."""
import contextlib
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_grad # pylint: disable=unused-import
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import rnn
from tensorflow.python.ops import rnn_cell
from tensorflow.python.ops import tensor_array_grad # pylint: disable=unused-import
from tensorflow.python.ops import variable_scope
from tensorflow.python.profiler import model_analyzer
from tensorflow.python.training import gradient_descent
from tensorflow.python.util import _pywrap_tfprof as print_mdl
from tensorflow.python.util import compat
def BuildSmallModel():
"""Build a small forward conv model."""
image = array_ops.zeros([2, 6, 6, 3])
_ = variable_scope.get_variable(
'ScalarW', [],
dtypes.float32,
initializer=init_ops.random_normal_initializer(stddev=0.001))
kernel = variable_scope.get_variable(
'DW', [3, 3, 3, 6],
dtypes.float32,
initializer=init_ops.random_normal_initializer(stddev=0.001))
x = nn_ops.conv2d(image, kernel, [1, 2, 2, 1], padding='SAME')
kernel = variable_scope.get_variable(
'DW2', [2, 2, 6, 12],
dtypes.float32,
initializer=init_ops.random_normal_initializer(stddev=0.001))
x = nn_ops.conv2d(x, kernel, [1, 2, 2, 1], padding='SAME')
return x
def BuildFullModel():
"""Build the full model with conv,rnn,opt."""
seq = []
for i in range(4):
with variable_scope.variable_scope('inp_%d' % i):
seq.append(array_ops.reshape(BuildSmallModel(), [2, 1, -1]))
cell = rnn_cell.BasicRNNCell(16)
out = rnn.dynamic_rnn(
cell, array_ops.concat(seq, axis=1), dtype=dtypes.float32)[0]
target = array_ops.ones_like(out)
loss = nn_ops.l2_loss(math_ops.reduce_mean(target - out))
sgd_op = gradient_descent.GradientDescentOptimizer(1e-2)
return sgd_op.minimize(loss)
def BuildSplittableModel():
"""Build a small model that can be run partially in each step."""
image = array_ops.zeros([2, 6, 6, 3])
kernel1 = variable_scope.get_variable(
'DW', [3, 3, 3, 6],
dtypes.float32,
initializer=init_ops.random_normal_initializer(stddev=0.001))
r1 = nn_ops.conv2d(image, kernel1, [1, 2, 2, 1], padding='SAME')
kernel2 = variable_scope.get_variable(
'DW2', [2, 3, 3, 6],
dtypes.float32,
initializer=init_ops.random_normal_initializer(stddev=0.001))
r2 = nn_ops.conv2d(image, kernel2, [1, 2, 2, 1], padding='SAME')
r3 = r1 + r2
return r1, r2, r3
def SearchTFProfNode(node, name):
"""Search a node in the tree."""
if node.name == name:
return node
for c in node.children:
r = SearchTFProfNode(c, name)
if r: return r
return None
@contextlib.contextmanager
def ProfilerFromFile(profile_file):
"""Initialize a profiler from profile file."""
print_mdl.ProfilerFromFile(compat.as_bytes(profile_file))
profiler = model_analyzer.Profiler.__new__(model_analyzer.Profiler)
yield profiler
print_mdl.DeleteProfiler()
def CheckAndRemoveDoc(profile):
assert 'Doc:' in profile
start_pos = profile.find('Profile:')
return profile[start_pos + 9:]
@@ -0,0 +1,61 @@
# Copyright 2016 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.
# ==============================================================================
"""print_model_analysis test."""
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.platform import test
# pylint: disable=bad-whitespace
# pylint: disable=bad-continuation
TEST_OPTIONS = {
'max_depth': 10000,
'min_bytes': 0,
'min_micros': 0,
'min_params': 0,
'min_float_ops': 0,
'order_by': 'name',
'account_type_regexes': ['.*'],
'start_name_regexes': ['.*'],
'trim_name_regexes': [],
'show_name_regexes': ['.*'],
'hide_name_regexes': [],
'account_displayed_op_only': True,
'select': ['params'],
'output': 'stdout',
}
# pylint: enable=bad-whitespace
# pylint: enable=bad-continuation
class PrintModelAnalysisTest(test.TestCase):
def _BuildSmallModel(self):
image = array_ops.zeros([2, 6, 6, 3])
kernel = variable_scope.get_variable(
'DW', [6, 6, 3, 6],
dtypes.float32,
initializer=init_ops.random_normal_initializer(stddev=0.001))
x = nn_ops.conv2d(image, kernel, [1, 2, 2, 1], padding='SAME')
return x
if __name__ == '__main__':
test.main()
@@ -0,0 +1,73 @@
/* 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.
==============================================================================*/
#include "tensorflow/python/profiler/internal/profiler_pywrap_impl.h"
#include <string>
#include <variant>
#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/profiler/convert/xplane_to_trace_events.h"
#include "xla/tsl/profiler/rpc/client/capture_profile.h"
#include "xla/tsl/profiler/utils/session_manager.h"
#include "tensorflow/core/platform/types.h"
#include "tsl/profiler/lib/profiler_session.h"
#include "tsl/profiler/protobuf/xplane.pb.h"
namespace tensorflow {
namespace profiler {
namespace pywrap {
using tsl::profiler::GetRemoteSessionManagerOptionsLocked;
absl::Status ProfilerSessionWrapper::Start(
const char* logdir,
const absl::flat_hash_map<std::string,
std::variant<bool, int, std::string>>& options) {
auto opts = GetRemoteSessionManagerOptionsLocked(logdir, options);
session_ = tsl::ProfilerSession::Create(opts.profiler_options());
logdir_ = logdir;
return session_->Status();
}
absl::Status ProfilerSessionWrapper::Stop(std::string* result) {
if (session_ != nullptr) {
tensorflow::profiler::XSpace xspace;
absl::Status status = session_->CollectData(&xspace);
session_.reset();
tsl::profiler::ConvertXSpaceToTraceEventsString(xspace, result);
TF_RETURN_IF_ERROR(status);
}
return absl::OkStatus();
}
absl::Status ProfilerSessionWrapper::ExportToTensorBoard() {
if (!session_ || logdir_.empty()) {
return absl::OkStatus();
}
tensorflow::profiler::XSpace xspace;
absl::Status status;
status = session_->CollectData(&xspace);
session_.reset();
status = tsl::profiler::ExportToTensorBoard(xspace, logdir_);
return status;
}
} // namespace pywrap
} // namespace profiler
} // namespace tensorflow
@@ -0,0 +1,50 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_PYTHON_PROFILER_INTERNAL_PROFILER_PYWRAP_IMPL_H_
#define TENSORFLOW_PYTHON_PROFILER_INTERNAL_PROFILER_PYWRAP_IMPL_H_
#include <memory>
#include <string>
#include <variant>
#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/types/variant.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/profiler/lib/profiler_session.h"
namespace tensorflow {
namespace profiler {
namespace pywrap {
class ProfilerSessionWrapper {
public:
absl::Status Start(
const char* logdir,
const absl::flat_hash_map<std::string,
std::variant<bool, int, std::string>>& options);
absl::Status Stop(std::string* result);
absl::Status ExportToTensorBoard();
private:
std::unique_ptr<tsl::ProfilerSession> session_;
std::string logdir_;
};
} // namespace pywrap
} // namespace profiler
} // namespace tensorflow
#endif // TENSORFLOW_PYTHON_PROFILER_INTERNAL_PROFILER_PYWRAP_IMPL_H_
@@ -0,0 +1,111 @@
/* Copyright 2020 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.
==============================================================================*/
#include <memory>
#include <string>
#include <utility>
#include <variant>
#include "absl/status/status.h"
#include "pybind11/pybind11.h" // from @pybind11
#include "tensorflow/core/profiler/rpc/profiler_server.h"
#include "tensorflow/python/lib/core/pybind11_status.h"
#include "tensorflow/python/profiler/internal/profiler_pywrap_impl.h"
#include "xprof/convert/repository.h" // from @org_xprof
#include "xprof/convert/tool_options.h" // from @org_xprof
#include "xprof/convert/xplane_to_tools_data.h" // from @org_xprof
namespace py = ::pybind11;
namespace {
using ::tensorflow::profiler::ToolOptions;
using ::tensorflow::profiler::pywrap::ProfilerSessionWrapper;
// These must be called under GIL because it reads Python objects. Reading
// Python objects require GIL because the objects can be mutated by other Python
// threads. In addition, Python objects are reference counted; reading py::dict
// will increase its reference count.
ToolOptions ToolOptionsFromPythonDict(const py::dict& dictionary) {
ToolOptions map;
for (const auto& item : dictionary) {
std::variant<bool, int, std::string> value;
try {
value = item.second.cast<bool>();
} catch (...) {
try {
value = item.second.cast<int>();
} catch (...) {
try {
value = item.second.cast<std::string>();
} catch (...) {
continue;
}
}
}
map.emplace(item.first.cast<std::string>(), value);
}
return map;
}
} // namespace
PYBIND11_MODULE(_pywrap_profiler, m) {
py::class_<ProfilerSessionWrapper> profiler_session_class(m,
"ProfilerSession");
profiler_session_class.def(py::init<>())
.def("start",
[](ProfilerSessionWrapper& wrapper, const char* logdir,
const py::dict& options) {
absl::Status status;
ToolOptions tool_options = ToolOptionsFromPythonDict(options);
{
py::gil_scoped_release release;
status = wrapper.Start(logdir, tool_options);
}
// Py_INCREF and Py_DECREF must be called holding the GIL.
tensorflow::MaybeRaiseRegisteredFromStatus(status);
})
.def("stop",
[](ProfilerSessionWrapper& wrapper) {
std::string content;
absl::Status status;
{
py::gil_scoped_release release;
status = wrapper.Stop(&content);
}
// Py_INCREF and Py_DECREF must be called holding the GIL.
tensorflow::MaybeRaiseRegisteredFromStatus(status);
// The content is not valid UTF-8. It must be converted to bytes.
return py::bytes(content);
})
.def("export_to_tb", [](ProfilerSessionWrapper& wrapper) {
absl::Status status;
{
py::gil_scoped_release release;
status = wrapper.ExportToTensorBoard();
}
// Py_INCREF and Py_DECREF must be called holding the GIL.
tensorflow::MaybeRaiseRegisteredFromStatus(status);
});
m.def("start_server", [](int port) {
auto profiler_server = std::make_unique<tsl::profiler::ProfilerServer>();
profiler_server->StartProfilerServer(port);
// Intentionally release profiler server. Should transfer ownership to
// caller instead.
profiler_server.release();
});
};
@@ -0,0 +1,48 @@
/* Copyright 2020 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.
==============================================================================*/
#ifndef TENSORFLOW_PYTHON_PROFILER_INTERNAL_PYTHON_HOOKS_H_
#define TENSORFLOW_PYTHON_PROFILER_INTERNAL_PYTHON_HOOKS_H_
#include <memory>
#include <stack>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/memory/memory.h"
#include "pybind11/cast.h" // from @pybind11
#include "pybind11/pybind11.h" // from @pybind11
#include "pybind11/pytypes.h" // from @pybind11
#include "xla/python/profiler/internal/python_hooks.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/types.h"
#include "tsl/profiler/protobuf/xplane.pb.h"
namespace tensorflow {
namespace profiler {
using xla::profiler::PythonHooksOptions; // NOLINT
using xla::profiler::PythonTraceEntry; // NOLINT
using xla::profiler::PerThreadEvents; // NOLINT
using xla::profiler::PythonHookContext; // NOLINT
using xla::profiler::PythonHooks; // NOLINT
} // namespace profiler
} // namespace tensorflow
#endif // TENSORFLOW_PYTHON_PROFILER_INTERNAL_PYTHON_HOOKS_H_
@@ -0,0 +1,175 @@
/* Copyright 2024 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.
==============================================================================*/
#include <string>
#include <utility>
#include <variant>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "pybind11/pybind11.h" // from @pybind11
#include "xla/pjrt/status_casters.h"
#include "xla/tsl/platform/types.h"
#include "xla/tsl/profiler/rpc/client/capture_profile.h"
#include "xprof/convert/tool_options.h" // from @org_xprof
#include "xprof/pywrap/profiler_plugin_impl.h" // from @org_xprof
namespace py = ::pybind11;
namespace {
using ::tensorflow::profiler::ToolOptions;
// These must be called under GIL because it reads Python objects. Reading
// Python objects require GIL because the objects can be mutated by other Python
// threads. In addition, Python objects are reference counted; reading py::dict
// will increase its reference count.
ToolOptions ToolOptionsFromPythonDict(const py::dict& dictionary) {
ToolOptions map;
// Copy keys to avoid iterator invalidation if the dict is mutated
// during implicit casts of values which can run Python user code.
std::vector<py::object> keys;
keys.reserve(dictionary.size());
for (const auto& item : dictionary) {
keys.push_back(py::reinterpret_borrow<py::object>(item.first));
}
for (const auto& key : keys) {
if (!dictionary.contains(key)) continue;
py::object py_value = dictionary[key];
std::variant<bool, int, std::string> value;
try {
value = py_value.cast<bool>();
} catch (...) {
try {
value = py_value.cast<int>();
} catch (...) {
try {
value = py_value.cast<std::string>();
} catch (...) {
continue;
}
}
}
try {
map.emplace(key.cast<std::string>(), value);
} catch (...) {
// Ignore keys that can't be cast to string
}
}
return map;
}
} // namespace
PYBIND11_MODULE(_pywrap_profiler_plugin, m, pybind11::mod_gil_not_used()) {
m.def(
"trace", [](const char* service_addr, const char* logdir,
const char* worker_list, bool include_dataset_ops,
int duration_ms, int num_tracing_attempts, py::dict options) {
absl::Status status;
ToolOptions tool_options = ToolOptionsFromPythonDict(options);
{
py::gil_scoped_release release;
status = tsl::profiler::CaptureRemoteTrace(
service_addr, logdir, worker_list, include_dataset_ops,
duration_ms, num_tracing_attempts, tool_options);
}
// Py_INCREF and Py_DECREF must be called holding the GIL.
xla::ThrowIfError(status);
});
m.def("monitor", [](const char* service_addr, int duration_ms,
int monitoring_level, bool display_timestamp) {
std::string content;
absl::Status status;
{
py::gil_scoped_release release;
status =
xprof::pywrap::Monitor(service_addr, duration_ms, monitoring_level,
display_timestamp, &content);
}
// Py_INCREF and Py_DECREF must be called holding the GIL.
xla::ThrowIfError(status);
return content;
});
m.def(
"xspace_to_tools_data",
[](const py::list& xspace_path_list, const py::str& py_tool_name,
const py::dict options = py::dict()) {
std::vector<std::string> xspace_paths;
xspace_paths.reserve(xspace_path_list.size());
for (py::handle obj : xspace_path_list) {
std::string xspace_path = std::string(py::cast<py::str>(obj));
xspace_paths.push_back(xspace_path);
}
std::string tool_name = std::string(py_tool_name);
ToolOptions tool_options = ToolOptionsFromPythonDict(options);
absl::StatusOr<std::pair<std::string, bool>> result;
{
py::gil_scoped_release release;
result = xprof::pywrap::XSpaceToToolsData(xspace_paths, tool_name,
tool_options);
}
if (!result.ok()) {
xla::ThrowIfError(result.status());
}
return py::make_tuple(py::bytes(result->first),
py::bool_(result->second));
},
// TODO: consider defaulting `xspace_path_list` to empty list, since
// this parameter is only used for two of the tools.
py::arg(), py::arg(), py::arg() = py::dict());
m.def(
"xspace_to_tools_data_from_byte_string",
[](const py::list& xspace_string_list, const py::list& filenames_list,
const py::str& py_tool_name, const py::dict options = py::dict()) {
std::vector<std::string> xspace_strings;
xspace_strings.reserve(xspace_string_list.size());
for (py::handle obj : xspace_string_list) {
xspace_strings.push_back(std::string(py::cast<py::bytes>(obj)));
}
std::vector<std::string> xspace_paths;
xspace_paths.reserve(filenames_list.size());
for (py::handle obj : filenames_list) {
xspace_paths.push_back(std::string(py::cast<py::str>(obj)));
}
std::string tool_name = std::string(py_tool_name);
ToolOptions tool_options = ToolOptionsFromPythonDict(options);
absl::StatusOr<std::pair<std::string, bool>> result;
{
py::gil_scoped_release release;
result = xprof::pywrap::XSpaceToToolsDataFromByteString(
xspace_strings, xspace_paths, tool_name, tool_options);
}
if (!result.ok()) {
xla::ThrowIfError(result.status());
}
return py::make_tuple(py::bytes(result->first),
py::bool_(result->second));
},
py::arg(), py::arg(), py::arg(), py::arg() = py::dict());
};
@@ -0,0 +1,242 @@
# Copyright 2016 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.
# ==============================================================================
"""test the RunMetadata proto."""
from collections import defaultdict
from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.python.client import session
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.profiler import option_builder
# pylint: disable=g-bad-import-order
# XXX: this depends on pywrap_tensorflow and must come later
from tensorflow.python.profiler import model_analyzer
from tensorflow.python.profiler.internal import model_analyzer_testlib as lib
SIZE = 1300
builder = option_builder.ProfileOptionBuilder
def _extract_node(run_meta, node_name):
ret = defaultdict(list)
for dev_stat in run_meta.step_stats.dev_stats:
dev = dev_stat.device.lower()
if dev.find('cpu:') > 0:
dev = dev[dev.find('cpu:'):]
elif dev.find('gpu:') > 0:
dev = dev[dev.find('gpu:'):]
elif '/host:cpu' not in dev:
assert False, 'Unrecognized device name: %s' % dev
for node_stat in dev_stat.node_stats:
nname = node_stat.node_name
if nname.find(':') > 0:
nname = nname[:nname.find(':')]
if nname == node_name:
ret[dev].append(node_stat)
return ret
def _run_model():
x = random_ops.random_normal(shape=[1, SIZE])
w = random_ops.random_normal(shape=[SIZE, 2 * SIZE])
y = math_ops.matmul(x, w)
config = config_pb2.ConfigProto()
config.graph_options.rewrite_options.arithmetic_optimization = (
rewriter_config_pb2.RewriterConfig.OFF)
with session.Session(config=config) as sess:
run_metadata = config_pb2.RunMetadata()
opts = builder.time_and_memory()
opts['min_micros'] = 0
opts['min_bytes'] = 0
opts['order_by'] = 'name'
opts['output'] = 'none'
_ = sess.run(
y,
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.SOFTWARE_TRACE),
run_metadata=run_metadata)
tfprof_node = model_analyzer.profile(
sess.graph, run_meta=run_metadata, options=opts)
return tfprof_node, run_metadata
def _run_loop_model():
config = config_pb2.ConfigProto()
# Grappler might fuse MatMul with BiasAdd in remapper optimizer.
config.graph_options.rewrite_options.remapping = (
rewriter_config_pb2.RewriterConfig.OFF)
with session.Session(config=config) as sess:
x = lib.BuildFullModel()
sess.run(variables.global_variables_initializer())
run_meta = config_pb2.RunMetadata()
_ = sess.run(
x,
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.SOFTWARE_TRACE),
run_metadata=run_meta)
opts = builder.time_and_memory()
opts['order_by'] = 'name'
opts['output'] = 'none'
tfprof_node = model_analyzer.profile(sess.graph, run_meta, options=opts)
return tfprof_node, run_meta
class RunMetadataTest(test.TestCase):
# This test requires HARDWARE_TRACE or FULL_TRACE to be specified to
# work as expected. Since we now run this test with SOFTWARE_TRACE
# (see _run_model routine above), this test will / should fail since
# GPU device tracers are not enabled
@test_util.run_deprecated_v1
def testGPU(self):
if not test.is_gpu_available(cuda_only=True):
return
gpu_dev = test.gpu_device_name()
ops.reset_default_graph()
with ops.device(gpu_dev):
tfprof_node, run_meta = _run_model()
self.assertEqual(tfprof_node.children[0].name, 'MatMul')
self.assertGreater(tfprof_node.children[0].exec_micros, 10)
ret = _extract_node(run_meta, 'MatMul')
self.assertEqual(len(ret['gpu:0']), 1)
@test_util.run_deprecated_v1
def testAllocationHistory(self):
if not test.is_gpu_available(cuda_only=True):
return
gpu_dev = test.gpu_device_name()
ops.reset_default_graph()
with ops.device(gpu_dev):
_, run_meta = _run_model()
mm = _extract_node(run_meta, 'MatMul')['gpu:0'][0]
mm_allocs = mm.memory[0].allocation_records
# has allocation and deallocation.
self.assertEqual(len(mm_allocs), 2)
# first allocated.
self.assertGreater(mm_allocs[1].alloc_micros, mm_allocs[0].alloc_micros)
self.assertGreater(mm_allocs[0].alloc_bytes, 0)
# Then deallocated.
self.assertLess(mm_allocs[1].alloc_bytes, 0)
# All memory deallocated.
self.assertEqual(mm_allocs[0].alloc_bytes + mm_allocs[1].alloc_bytes, 0)
rand = _extract_node(run_meta,
'random_normal/RandomStandardNormal')['gpu:0'][0]
random_allocs = rand.memory[0].allocation_records
# random normal must allocated first since matmul depends on it.
self.assertLess(random_allocs[0].alloc_micros, mm.all_start_micros)
# deallocates the memory after matmul started.
self.assertGreater(random_allocs[1].alloc_micros, mm.all_start_micros)
@test_util.run_deprecated_v1
def testCPU(self):
ops.reset_default_graph()
with ops.device('/cpu:0'):
tfprof_node, run_meta = _run_model()
self.assertEqual(tfprof_node.children[0].name, 'MatMul')
self.assertGreater(tfprof_node.children[0].exec_micros, 0)
ret = _extract_node(run_meta, 'MatMul')
self.assertEqual(len(ret['cpu:0']), 1)
ret = _extract_node(run_meta, 'MatMul:MatMul')
self.assertEqual(len(ret), 0)
@test_util.run_v1_only('b/120545219')
def testLoopCPU(self):
ops.reset_default_graph()
with ops.device('/cpu:0'):
tfprof_node, run_meta = _run_loop_model()
# The while-loop caused a node to appear 4 times in scheduling.
ret = _extract_node(run_meta, 'rnn/while/basic_rnn_cell/MatMul')
self.assertEqual(len(ret['cpu:0']), 4)
total_cpu_execs = 0
for node in ret['cpu:0']:
total_cpu_execs += node.op_end_rel_micros
mm_node = lib.SearchTFProfNode(tfprof_node,
'rnn/while/basic_rnn_cell/MatMul')
self.assertEqual(mm_node.run_count, 4)
self.assertEqual(mm_node.cpu_exec_micros, total_cpu_execs)
self.assertEqual(mm_node.exec_micros, total_cpu_execs)
def testGradientGraph(self):
# Note: Please don't just adjust the test to make it pass.
# The code view logic depends on it.
ops.reset_default_graph()
_, _ = _run_loop_model()
graph = ops.get_default_graph()
forward_op = set()
backward_op = set()
back_to_forward = {}
for op in graph.get_operations():
if op.name.find('gradients/') > 0 and op.name.find('_grad/') > 0:
backward_op.add(op.name)
idx1 = op.name.find('gradients/') + 10
idx2 = op.name.find('_grad/')
back_to_forward[op.name] = op.name[idx1:idx2]
else:
forward_op.add(op.name)
for _, f in back_to_forward.items():
self.assertTrue(f in forward_op)
# This test requires HARDWARE_TRACE or FULL_TRACE to be specified to
# work as expected. Since we now run this test with SOFTWARE_TRACE
# (see _run_model routine above), this test will / should fail since
# GPU device tracers are not enabled
@test.disable_with_predicate(
pred=test.is_built_with_rocm,
skip_message='Test fails on ROCm when run without FULL_TRACE')
def testLoopGPU(self):
if not test.is_gpu_available():
return
ops.reset_default_graph()
with ops.device('/device:GPU:0'):
_, run_meta = _run_loop_model()
# The while-loop caused a node to appear 4 times in scheduling.
ret = _extract_node(run_meta, 'rnn/while/basic_rnn_cell/MatMul')
self.assertEqual(len(ret['gpu:0']), 4, '%s' % run_meta)
total_cpu_execs = 0
for node in ret['gpu:0']:
total_cpu_execs += node.op_end_rel_micros
self.assertGreaterEqual(
len(ret['gpu:0/stream:all']), 4, '%s' % run_meta)
if __name__ == '__main__':
test.main()
@@ -0,0 +1,49 @@
/* 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.
==============================================================================*/
#include "xla/python/profiler/internal/traceme_wrapper.h"
#include "pybind11/attr.h" // from @pybind11
#include "pybind11/pybind11.h" // from @pybind11
#include "xla/python/profiler/internal/traceme_state.h"
namespace py = ::pybind11;
using ::xla::profiler::TraceMeWrapper;
// Returns true if TraceMe is enabled.
// This is a low-overhead function that can be called frequently.
static PyObject* traceme_enabled(PyObject* self, PyObject* args) {
if (xla::profiler::traceme_enabled) {
Py_RETURN_TRUE;
}
Py_RETURN_FALSE;
}
static PyMethodDef traceme_method_def = {"traceme_enabled", traceme_enabled,
METH_NOARGS,
"Returns true if TraceMe is enabled."};
PYBIND11_MODULE(_pywrap_traceme, m) {
py::class_<TraceMeWrapper>(m, "TraceMe", py::module_local())
.def(py::init<const py::str&, const py::kwargs&>())
.def("SetMetadata", &TraceMeWrapper::SetMetadata)
.def("Stop", &TraceMeWrapper::Stop);
py::object module_name = m.attr("__name__");
m.attr("traceme_enabled") =
py::reinterpret_steal<py::object>(PyCFunction_NewEx(
&traceme_method_def, /*self=*/nullptr, module_name.ptr()));
};
@@ -0,0 +1,422 @@
# Copyright 2016 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.
# ==============================================================================
"""Model Analyzer.
Analyze model, including shape, params, time, memory, structure, etc.
"""
import sys
from google.protobuf import message
from tensorflow.core.profiler import tfprof_options_pb2
from tensorflow.core.profiler import tfprof_output_pb2
from tensorflow.python.eager import context
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.profiler import option_builder
from tensorflow.python.profiler import tfprof_logger
from tensorflow.python.util import _pywrap_tfprof as print_mdl
from tensorflow.python.util.tf_export import tf_export
_DEFAULT_PROFILE_OPTIONS = 0
_DEFAULT_ADVISE_OPTIONS = 0
# The following options are for 'advise' cmd.
# Show all advice.
ALL_ADVICE = {
'ExpensiveOperationChecker': {},
'AcceleratorUtilizationChecker': {},
'JobChecker': {}, # Only available internally.
'OperationChecker': {},
}
def _graph_string(graph):
"""Helper to serialize a graph to string."""
if graph:
return graph.as_graph_def(add_shapes=True).SerializeToString()
else:
return b''
def _build_options(options):
"""Build tfprof.OptionsProto.
Args:
options: A dictionary of options.
Returns:
tfprof.OptionsProto.
"""
opts = tfprof_options_pb2.OptionsProto()
opts.max_depth = options.get('max_depth', 10)
opts.min_bytes = options.get('min_bytes', 0)
opts.min_peak_bytes = options.get('min_peak_bytes', 0)
opts.min_residual_bytes = options.get('min_residual_bytes', 0)
opts.min_output_bytes = options.get('min_output_bytes', 0)
opts.min_micros = options.get('min_micros', 0)
opts.min_accelerator_micros = options.get('min_accelerator_micros', 0)
opts.min_cpu_micros = options.get('min_cpu_micros', 0)
opts.min_params = options.get('min_params', 0)
opts.min_float_ops = options.get('min_float_ops', 0)
opts.min_occurrence = options.get('min_occurrence', 0)
opts.step = options.get('step', -1)
opts.order_by = options.get('order_by', 'name')
for p in options.get('account_type_regexes', []):
opts.account_type_regexes.append(p)
for p in options.get('start_name_regexes', []):
opts.start_name_regexes.append(p)
for p in options.get('trim_name_regexes', []):
opts.trim_name_regexes.append(p)
for p in options.get('show_name_regexes', []):
opts.show_name_regexes.append(p)
for p in options.get('hide_name_regexes', []):
opts.hide_name_regexes.append(p)
opts.account_displayed_op_only = options.get('account_displayed_op_only',
False)
for p in options.get('select', []):
opts.select.append(p)
opts.output = options.get('output', 'stdout')
opts.dump_to_file = options.get('dump_to_file', '')
return opts
def _build_advisor_options(options):
"""Build tfprof.AdvisorOptionsProto.
Args:
options: A dictionary of options. See ALL_ADVICE example.
Returns:
tfprof.AdvisorOptionsProto.
"""
opts = tfprof_options_pb2.AdvisorOptionsProto()
if options is None:
return opts
for checker, checker_opts in options.items():
checker_ops_pb = tfprof_options_pb2.AdvisorOptionsProto.CheckerOption()
for k, v in checker_opts.items():
checker_ops_pb[k] = v
opts.checkers[checker].MergeFrom(checker_ops_pb)
return opts
@tf_export(v1=['profiler.Profiler'])
class Profiler:
"""TensorFlow multi-step profiler.
```python
Typical use case:
# Currently we are only allowed to create 1 profiler per process.
profiler = Profiler(sess.graph)
for i in range(total_steps):
if i % 10000 == 0:
run_meta = tf.compat.v1.RunMetadata()
_ = sess.run(...,
options=tf.compat.v1.RunOptions(
trace_level=tf.RunOptions.FULL_TRACE),
run_metadata=run_meta)
profiler.add_step(i, run_meta)
# Profile the parameters of your model.
profiler.profile_name_scope(options=(option_builder.ProfileOptionBuilder
.trainable_variables_parameter()))
# Or profile the timing of your model operations.
opts = option_builder.ProfileOptionBuilder.time_and_memory()
profiler.profile_operations(options=opts)
# Or you can generate a timeline:
opts = (option_builder.ProfileOptionBuilder(
option_builder.ProfileOptionBuilder.time_and_memory())
.with_step(i)
.with_timeline_output(filename).build())
profiler.profile_graph(options=opts)
else:
_ = sess.run(...)
# Auto detect problems and generate advice.
profiler.advise()
```
"""
def __init__(self, graph=None, op_log=None):
"""Constructor.
Args:
graph: tf.Graph. If None and eager execution is not enabled, use default
graph.
op_log: optional. tensorflow::tfprof::OpLogProto proto. Used to define
extra op types.
"""
if not graph and not context.executing_eagerly():
graph = ops.get_default_graph()
self._coverage = 0.0
self._graph = graph
# pylint: disable=protected-access
op_log = tfprof_logger.merge_default_with_oplog(self._graph, op_log=op_log)
# pylint: enable=protected-access
print_mdl.NewProfiler(
_graph_string(self._graph), op_log.SerializeToString())
def __del__(self):
print_mdl.DeleteProfiler()
def add_step(self, step, run_meta):
"""Add statistics of a step.
Args:
step: int, An id used to group one or more different `run_meta` together.
When profiling with the profile_xxx APIs, user can use the `step` id in
the `options` to profile these `run_meta` together.
run_meta: RunMetadata proto that contains statistics of a session run.
"""
# pylint: disable=protected-access
op_log = tfprof_logger.merge_default_with_oplog(
self._graph, run_meta=run_meta)
# pylint: enable=protected-access
# TODO(xpan): P1: Better to find the current graph.
self._coverage = print_mdl.AddStep(step, _graph_string(self._graph),
run_meta.SerializeToString(),
op_log.SerializeToString())
def profile_python(self, options):
"""Profile the statistics of the Python codes.
By default, it shows the call stack from root. To avoid
redundant output, you may use options to filter as below
options['show_name_regexes'] = ['.*my_code.py.*']
Args:
options: A dict of options. See core/profiler/g3doc/options.md.
Returns:
a MultiGraphNodeProto that records the results.
"""
opts = _build_options(options)
tfprof_node = tfprof_output_pb2.MultiGraphNodeProto()
try:
tfprof_node.ParseFromString(
print_mdl.Profile('code'.encode('utf-8'), opts.SerializeToString()))
except message.DecodeError as e:
sys.stderr.write('Cannot parse returned proto: %s.\n' % e)
return tfprof_node
def profile_operations(self, options):
"""Profile the statistics of the Operation types (e.g.
MatMul, Conv2D).
Args:
options: A dict of options. See core/profiler/g3doc/options.md.
Returns:
a MultiGraphNodeProto that records the results.
"""
opts = _build_options(options)
tfprof_node = tfprof_output_pb2.MultiGraphNodeProto()
try:
tfprof_node.ParseFromString(
print_mdl.Profile('op'.encode('utf-8'), opts.SerializeToString()))
except message.DecodeError as e:
sys.stderr.write('Cannot parse returned proto: %s.\n' % e)
return tfprof_node
def profile_name_scope(self, options):
"""Profile the statistics of graph nodes, organized by name scope.
Args:
options: A dict of options. See core/profiler/g3doc/options.md.
Returns:
a GraphNodeProto that records the results.
"""
opts = _build_options(options)
tfprof_node = tfprof_output_pb2.GraphNodeProto()
try:
tfprof_node.ParseFromString(
print_mdl.Profile('scope'.encode('utf-8'), opts.SerializeToString()))
except message.DecodeError as e:
sys.stderr.write('Cannot parse returned proto: %s.\n' % e)
return tfprof_node
def profile_graph(self, options):
"""Profile the statistics of graph nodes, organized by dataflow graph.
Args:
options: A dict of options. See core/profiler/g3doc/options.md.
Returns:
a GraphNodeProto that records the results.
"""
opts = _build_options(options)
tfprof_node = tfprof_output_pb2.GraphNodeProto()
try:
tfprof_node.ParseFromString(
print_mdl.Profile('graph'.encode('utf-8'), opts.SerializeToString()))
except message.DecodeError as e:
sys.stderr.write('Cannot parse returned proto: %s.\n' % e)
return tfprof_node
def advise(self, options):
"""Automatically detect problems and generate reports.
Args:
options: A dict of options. See ALL_ADVICE example above.
Returns:
An Advise proto that contains the reports from all checkers.
"""
advise_pb = tfprof_output_pb2.AdviceProto()
opts = _build_advisor_options(options)
advise_pb.ParseFromString(
print_mdl.Profile('advise'.encode('utf-8'), opts.SerializeToString()))
return advise_pb
def serialize_to_string(self):
"""Serialize the ProfileProto to a binary string.
Users can write it to file for offline analysis by tfprof commandline
or graphical interface.
Returns:
ProfileProto binary string.
"""
return print_mdl.SerializeToString()
def _write_profile(self, filename):
"""Writes the profile to a file."""
print_mdl.WriteProfile(filename)
@tf_export(v1=['profiler.profile'])
def profile(graph=None,
run_meta=None,
op_log=None,
cmd='scope',
options=_DEFAULT_PROFILE_OPTIONS):
"""Profile model.
Tutorials and examples can be found in:
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/profiler/g3doc/python_api.md
Args:
graph: tf.Graph. If None and eager execution is not enabled, use default
graph.
run_meta: optional tensorflow.RunMetadata proto. It is necessary to
support run time information profiling, such as time and memory.
op_log: tensorflow.tfprof.OpLogProto proto. User can assign "types" to graph
nodes with op_log. "types" allow user to flexibly group and account
profiles using options['accounted_type_regexes'].
cmd: string. Either 'op', 'scope', 'graph' or 'code'. 'op' view organizes
profile using operation type. (e.g. MatMul) 'scope' view organizes profile
using graph node name scope. 'graph' view organizes profile using graph
node inputs/outputs. 'code' view organizes profile using Python call
stack.
options: A dict of options. See core/profiler/g3doc/options.md.
Returns:
If cmd is 'scope' or 'graph', returns GraphNodeProto proto.
If cmd is 'op' or 'code', returns MultiGraphNodeProto proto.
Side effect: stdout/file/timeline.json depending on options['output']
"""
if not graph and not context.executing_eagerly():
graph = ops.get_default_graph()
if options == _DEFAULT_PROFILE_OPTIONS:
options = (
option_builder.ProfileOptionBuilder.trainable_variables_parameter())
# pylint: disable=protected-access
op_log = tfprof_logger.merge_default_with_oplog(
graph, op_log, run_meta, add_trace=cmd == 'code')
# pylint: enable=protected-access
opts = _build_options(options)
run_meta_str = run_meta.SerializeToString() if run_meta else b''
graph_str = _graph_string(graph)
if cmd == 'code' or cmd == 'op':
tfprof_node = tfprof_output_pb2.MultiGraphNodeProto()
ret = print_mdl.PrintModelAnalysis(graph_str, run_meta_str,
op_log.SerializeToString(),
cmd.encode('utf-8'),
opts.SerializeToString())
try:
tfprof_node.ParseFromString(ret)
except message.DecodeError as e:
sys.stderr.write('Cannot parse returned proto: %s.\n' % e)
elif cmd == 'graph' or cmd == 'scope':
tfprof_node = tfprof_output_pb2.GraphNodeProto()
ret = print_mdl.PrintModelAnalysis(graph_str, run_meta_str,
op_log.SerializeToString(),
cmd.encode('utf-8'),
opts.SerializeToString())
try:
tfprof_node.ParseFromString(ret)
except message.DecodeError as e:
sys.stderr.write('Cannot parse returned proto: %s.\n' % e)
else:
raise errors.InvalidArgumentError(None, None, 'unknown cmd: %s\n' % cmd)
return tfprof_node
@tf_export(v1=['profiler.advise'])
def advise(graph=None, run_meta=None, options=_DEFAULT_ADVISE_OPTIONS):
"""Auto profile and advise.
Builds profiles and automatically check anomalies of various
aspects. For more details:
https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core/profiler/README.md
Args:
graph: tf.Graph. If None and eager execution is not enabled, use default
graph.
run_meta: optional tensorflow.RunMetadata proto. It is necessary to
support run time information profiling, such as time and memory.
options: see ALL_ADVICE example above. Default checks everything.
Returns:
Returns AdviceProto proto
"""
if not graph and not context.executing_eagerly():
graph = ops.get_default_graph()
if options == _DEFAULT_ADVISE_OPTIONS:
options = ALL_ADVICE.copy()
# pylint: disable=protected-access
op_log = tfprof_logger.merge_default_with_oplog(
graph, None, run_meta, add_trace=True)
# pylint: enable=protected-access
run_meta_str = run_meta.SerializeToString() if run_meta else b''
opts = _build_advisor_options(options)
ret = tfprof_output_pb2.AdviceProto()
ret.ParseFromString(
print_mdl.PrintModelAnalysis(
_graph_string(graph), run_meta_str, op_log.SerializeToString(),
'advise'.encode('utf-8'), opts.SerializeToString()))
return ret
@@ -0,0 +1,782 @@
# Copyright 2016 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.
# ==============================================================================
import gzip
import io
import os
import random
import re
import numpy as np
from tensorflow.core.profiler import profile_pb2
from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.python.client import session
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gradients
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import variables
from tensorflow.python.ops import while_loop
from tensorflow.python.platform import gfile
from tensorflow.python.platform import test
from tensorflow.python.profiler import model_analyzer
from tensorflow.python.profiler import option_builder
from tensorflow.python.profiler import profile_context
from tensorflow.python.profiler.internal import model_analyzer_testlib as lib
from tensorflow.python.util import compat
builder = option_builder.ProfileOptionBuilder
class PrintModelAnalysisTest(test.TestCase):
def _no_rewrite_session_config(self):
rewriter_config = rewriter_config_pb2.RewriterConfig(
pin_to_host_optimization=rewriter_config_pb2.RewriterConfig.OFF)
graph_options = config_pb2.GraphOptions(rewrite_options=rewriter_config)
return config_pb2.ConfigProto(graph_options=graph_options)
def testDumpToFile(self):
ops.reset_default_graph()
outfile = os.path.join(test.get_temp_dir(), 'dump')
opts = builder(builder.trainable_variables_parameter()).with_file_output(
outfile).build()
with session.Session(config=self._no_rewrite_session_config()) as sess:
_ = lib.BuildSmallModel()
model_analyzer.profile(sess.graph, options=opts)
with gfile.Open(outfile, 'r') as f:
self.assertEqual(
u'node name | # parameters\n'
'_TFProfRoot (--/451 params)\n'
' DW (3x3x3x6, 162/162 params)\n'
' DW2 (2x2x6x12, 288/288 params)\n'
' ScalarW (1, 1/1 params)\n', lib.CheckAndRemoveDoc(f.read()))
@test_util.run_v1_only('b/120545219')
def testSelectEverythingDetail(self):
ops.reset_default_graph()
dev = '/device:GPU:0' if test.is_gpu_available() else '/device:CPU:0'
outfile = os.path.join(test.get_temp_dir(), 'dump')
opts = (
builder(builder.trainable_variables_parameter()).with_file_output(
outfile).with_accounted_types(['.*']).select([
'micros', 'bytes', 'params', 'float_ops', 'occurrence',
'device', 'op_types', 'input_shapes'
]).build())
with profile_context.ProfileContext(
test.get_temp_dir(), trace_steps=[], dump_steps=[]) as pctx:
with session.Session(
config=self._no_rewrite_session_config()) as sess, ops.device(dev):
x = lib.BuildSmallModel()
self.evaluate(variables.global_variables_initializer())
pctx.trace_next_step()
pctx.dump_next_step()
_ = self.evaluate(x)
pctx.profiler.profile_name_scope(options=opts)
with gfile.Open(outfile, 'r') as f:
# pylint: disable=line-too-long
dump_str = lib.CheckAndRemoveDoc(f.read())
outputs = dump_str.split('\n')
self.assertEqual(
outputs[0],
'node name | # parameters | # float_ops | requested bytes | total execution time | accelerator execution time | cpu execution time | assigned devices | op types | op count (run|defined) | input shapes'
)
for o in outputs[1:]:
if o.find('Conv2D ') > 0:
metrics = o[o.find('(') + 1:o.find(')')].split(',')
# Make sure time is profiled.
gap = 1 if test.is_gpu_available() else 2
for i in range(3, 6, gap):
mat = re.search('(.*)(?:us|ms|sec)/(.*)(?:us|ms|sec)',
metrics[i])
self.assertGreater(float(mat.group(1)), 0.0)
self.assertGreater(float(mat.group(2)), 0.0)
# Make sure device is profiled.
if test.is_gpu_available():
self.assertTrue(metrics[6].find('gpu') > 0)
self.assertFalse(metrics[6].find('cpu') > 0)
else:
self.assertFalse(metrics[6].find('gpu') > 0)
self.assertTrue(metrics[6].find('cpu') > 0)
# Make sure float_ops is profiled.
mat = re.search('(.*)k/(.*)k flops', metrics[1].strip())
self.assertGreater(float(mat.group(1)), 0.0)
self.assertGreater(float(mat.group(2)), 0.0)
# Make sure op_count is profiled.
self.assertEqual(metrics[8].strip(), '1/1|1/1')
# Make sure input_shapes is profiled.
self.assertEqual(metrics[9].strip(), '0:2x6x6x3|1:3x3x3x6')
if o.find('DW (3x3x3x6') > 0:
metrics = o[o.find('(') + 1:o.find(')')].split(',')
mat = re.search('(.*)/(.*) params', metrics[1].strip())
self.assertGreater(float(mat.group(1)), 0.0)
self.assertGreater(float(mat.group(2)), 0.0)
# pylint: enable=line-too-long
# Test that profiler restored from profile file gives the same result.
gfile.Remove(outfile)
profile_file = os.path.join(test.get_temp_dir(), 'profile_1')
with lib.ProfilerFromFile(profile_file) as profiler:
profiler.profile_name_scope(options=opts)
with gfile.Open(outfile, 'r') as f:
self.assertEqual(dump_str, lib.CheckAndRemoveDoc(f.read()))
def testSelectEverything(self):
ops.reset_default_graph()
outfile = os.path.join(test.get_temp_dir(), 'dump')
opts = (
builder(builder.trainable_variables_parameter()).with_file_output(
outfile).with_accounted_types(['.*']).select([
'params', 'float_ops', 'occurrence', 'device', 'op_types',
'input_shapes'
]).build())
with session.Session(config=self._no_rewrite_session_config()
) as sess, ops.device('/device:CPU:0'):
x = lib.BuildSmallModel()
self.evaluate(variables.global_variables_initializer())
run_meta = config_pb2.RunMetadata()
_ = sess.run(
x,
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE),
run_metadata=run_meta)
model_analyzer.profile(sess.graph, run_meta, options=opts)
def testSimpleCodeView(self):
ops.reset_default_graph()
outfile = os.path.join(test.get_temp_dir(), 'dump')
# TODO(xpan): Test 'micros'. Since the execution time changes each run,
# it's a bit difficult to test it now.
opts = (
builder(builder.trainable_variables_parameter()).with_file_output(
outfile).with_accounted_types(['.*']).with_node_names(
show_name_regexes=['.*model_analyzer_testlib.*'
]).account_displayed_op_only(False).select([
'bytes', 'params', 'float_ops',
'num_hidden_ops', 'device', 'input_shapes'
]).build())
with session.Session(config=self._no_rewrite_session_config()) as sess:
x = lib.BuildSmallModel()
self.evaluate(variables.global_variables_initializer())
run_meta = config_pb2.RunMetadata()
_ = sess.run(
x,
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE),
run_metadata=run_meta)
model_analyzer.profile(sess.graph, run_meta, cmd='code', options=opts)
with gfile.Open(outfile, 'r') as f:
# pylint: disable=line-too-long
self.assertEqual(
'node name | requested bytes | # parameters | # float_ops | assigned devices | in',
lib.CheckAndRemoveDoc(f.read())[0:80])
# pylint: enable=line-too-long
@test_util.run_v1_only('b/120545219')
def testComplexCodeView(self):
ops.reset_default_graph()
outfile = os.path.join(test.get_temp_dir(), 'dump')
opts = (
builder(builder.trainable_variables_parameter()).with_file_output(
outfile).with_accounted_types(['.*']).with_node_names(
show_name_regexes=['.*model_analyzer_testlib.py.*'
]).account_displayed_op_only(False).select(
['params', 'float_ops']).build())
with profile_context.ProfileContext(
test.get_temp_dir(), trace_steps=[], dump_steps=[]) as pctx:
with session.Session(config=self._no_rewrite_session_config()) as sess:
x = lib.BuildFullModel()
self.evaluate(variables.global_variables_initializer())
pctx.trace_next_step()
_ = self.evaluate(x)
tfprof_node = pctx.profiler.profile_python(options=opts)
# pylint: disable=line-too-long
with gfile.Open(outfile, 'r') as f:
lines = f.read().split('\n')
self.assertGreater(len(lines), 5)
result = '\n'.join(l[:min(len(l), 80)] for l in lines)
self.assertTrue(
compat.as_text(lib.CheckAndRemoveDoc(result)).startswith(
'node name | # parameters | # float_ops'))
self.assertLess(0, tfprof_node.total_exec_micros)
self.assertEqual(2844, tfprof_node.total_parameters)
#The graph is modified when MKL is enabled,total_float_ops will
#be different
if test_util.IsMklEnabled():
self.assertLess(101600, tfprof_node.total_float_ops)
else:
self.assertLess(145660, tfprof_node.total_float_ops)
self.assertEqual(10, len(tfprof_node.children))
self.assertEqual('_TFProfRoot', tfprof_node.name)
self.assertEqual('model_analyzer_testlib.py:59:BuildFullModel',
tfprof_node.children[0].name)
self.assertEqual(
'model_analyzer_testlib.py:59:BuildFullModel (gradient)',
tfprof_node.children[1].name)
self.assertEqual('model_analyzer_testlib.py:62:BuildFullModel',
tfprof_node.children[2].name)
self.assertEqual(
'model_analyzer_testlib.py:62:BuildFullModel (gradient)',
tfprof_node.children[3].name)
self.assertEqual('model_analyzer_testlib.py:63:BuildFullModel',
tfprof_node.children[4].name)
self.assertEqual('model_analyzer_testlib.py:63:BuildFullModel (gradient)',
tfprof_node.children[5].name)
self.assertEqual(
'model_analyzer_testlib.py:65:BuildFullModel',
tfprof_node.children[6].name)
self.assertEqual('model_analyzer_testlib.py:66:BuildFullModel',
tfprof_node.children[7].name)
# pylint: enable=line-too-long
def testCodeViewLeafGraphNode(self):
ops.reset_default_graph()
opts = (
builder(builder.trainable_variables_parameter()).with_empty_output()
.with_accounted_types(['.*']).account_displayed_op_only(False).select(
['bytes', 'params', 'float_ops', 'device']).build())
with session.Session(config=self._no_rewrite_session_config()) as sess:
x = lib.BuildSmallModel()
self.evaluate(variables.global_variables_initializer())
run_meta = config_pb2.RunMetadata()
_ = sess.run(
x,
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE),
run_metadata=run_meta)
tfprof_node = model_analyzer.profile(
sess.graph, run_meta, cmd='code', options=opts)
leaf = tfprof_node
while leaf.children:
self.assertEqual(0, len(leaf.graph_nodes))
leaf = leaf.children[0]
self.assertEqual(1, len(leaf.graph_nodes))
def testTimeline(self):
ops.reset_default_graph()
outfile = os.path.join(test.get_temp_dir(), 'timeline')
opts = (
builder(builder.trainable_variables_parameter()).with_max_depth(100000)
.with_step(0).with_timeline_output(outfile).with_accounted_types(
['.*']).build())
with session.Session(config=self._no_rewrite_session_config()) as sess:
x = lib.BuildFullModel()
self.evaluate(variables.global_variables_initializer())
run_meta = config_pb2.RunMetadata()
_ = sess.run(
x,
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE),
run_metadata=run_meta)
_ = model_analyzer.profile(
sess.graph, run_meta, cmd='graph', options=opts)
with gfile.Open(outfile + '_0', 'r') as f:
# Test that a json file is created.
# TODO(xpan): tfprof Timeline isn't quite correct on Windows.
# Investigate why.
if os.name != 'nt':
self.assertLess(1000, len(f.read()))
else:
self.assertLess(1, len(f.read()))
def testOpView(self):
ops.reset_default_graph()
outfile = os.path.join(test.get_temp_dir(), 'dump')
opts = (
builder(builder.trainable_variables_parameter()).with_file_output(
outfile).with_accounted_types(
['.*']).with_min_occurrence(10).order_by('occurrence').select([
'params', 'micros', 'bytes', 'peak_bytes', 'residual_bytes',
'output_bytes', 'occurrence', 'input_shapes'
]).build())
with session.Session(config=self._no_rewrite_session_config()) as sess:
x = lib.BuildFullModel()
self.evaluate(variables.global_variables_initializer())
run_meta = config_pb2.RunMetadata()
_ = sess.run(
x,
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE),
run_metadata=run_meta)
tfprof_node = model_analyzer.profile(
sess.graph, run_meta, cmd='op', options=opts)
with gfile.Open(outfile, 'r') as f:
# pylint: disable=line-too-long
self.assertEqual(
'nodename|requestedbytes|peakbytes|residualbytes|outputbytes|totalexecutiontime|acceleratorexecutiontime|cpuexecutiontime|#parameters|opoccurrence(run|defined)|inputshapes',
lib.CheckAndRemoveDoc(f.read()).replace('\t',
'').replace(' ', '')[0:170])
# pylint: enable=line-too-long
total_children = 0
last_occurrence = 1e32
input_shapes = 0
last_total_micros = tfprof_node.total_exec_micros
last_micros = tfprof_node.exec_micros
while tfprof_node.children:
for gnode in tfprof_node.graph_nodes:
input_shapes += len(gnode.input_shapes)
self.assertEqual(len(tfprof_node.children), 1)
tfprof_node = tfprof_node.children[0]
self.assertEqual(last_total_micros,
tfprof_node.total_exec_micros + last_micros)
last_total_micros = tfprof_node.total_exec_micros
last_micros = tfprof_node.exec_micros
total_children += 1
self.assertLessEqual(len(tfprof_node.graph_nodes), last_occurrence)
last_occurrence = len(tfprof_node.graph_nodes)
self.assertGreater(input_shapes, 0)
def testAdvisor(self):
ops.reset_default_graph()
with session.Session(config=self._no_rewrite_session_config()) as sess:
x = lib.BuildFullModel()
self.evaluate(variables.global_variables_initializer())
run_meta = config_pb2.RunMetadata()
_ = sess.run(
x,
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE),
run_metadata=run_meta)
advice_pb = model_analyzer.advise(sess.graph, run_meta)
self.assertTrue('AcceleratorUtilizationChecker' in advice_pb.checkers)
self.assertTrue('ExpensiveOperationChecker' in advice_pb.checkers)
self.assertTrue('OperationChecker' in advice_pb.checkers)
checker = advice_pb.checkers['AcceleratorUtilizationChecker']
if test.is_gpu_available():
self.assertGreater(len(checker.reports), 0)
else:
self.assertEqual(len(checker.reports), 0)
checker = advice_pb.checkers['ExpensiveOperationChecker']
self.assertGreater(len(checker.reports), 0)
def pprof_test_helper(self, attribute, should_fail=False):
ops.reset_default_graph()
outfile = os.path.join(test.get_temp_dir(), attribute + '_pprof.pb.gz')
opts = (
builder(builder.time_and_memory()).select([
attribute
]).with_max_depth(100000).with_node_names(
trim_name_regexes=['ops.py.*']).with_pprof_output(outfile).build())
with session.Session(config=self._no_rewrite_session_config()) as sess:
x = lib.BuildFullModel()
self.evaluate(variables.global_variables_initializer())
run_meta = config_pb2.RunMetadata()
_ = sess.run(
x,
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE),
run_metadata=run_meta)
_ = model_analyzer.profile(sess.graph, run_meta, cmd='code', options=opts)
if should_fail:
self.assertFalse(gfile.Exists(outfile))
return
profile_pb = profile_pb2.Profile()
with gfile.Open(outfile, 'rb') as f:
with gzip.GzipFile(fileobj=io.BytesIO(f.read())) as gzipf:
profile_pb.ParseFromString(gzipf.read())
self.assertGreater(len(profile_pb.sample), 10)
self.assertGreater(len(profile_pb.location), 10)
self.assertGreater(len(profile_pb.function), 10)
self.assertGreater(len(profile_pb.string_table), 30)
has_rnn = False
for s in profile_pb.string_table:
if s.find('rnn') > 0:
has_rnn = True
self.assertFalse(s.startswith('ops.py'))
self.assertTrue(has_rnn)
def testPprof(self):
for attr in [
'micros', 'bytes', 'accelerator_micros', 'cpu_micros', 'params',
'float_ops'
]:
self.pprof_test_helper(attr)
for attr in ['op_types', 'device', 'input_shapes']:
self.pprof_test_helper(attr, True)
def testMinOption(self):
ops.reset_default_graph()
def check_min(nodes, mm=0, mam=0, mcm=0, mb=0, mpb=0, mrb=0, mob=0):
for n in nodes:
if mm > 0:
self.assertGreaterEqual(n.exec_micros, mm)
if mam > 0:
self.assertGreaterEqual(n.accelerator_exec_micros, mam)
if mcm > 0:
self.assertGreaterEqual(n.cpu_exec_micros, mcm)
if mb > 0:
self.assertGreaterEqual(n.requested_bytes, mb)
if mpb > 0:
self.assertGreaterEqual(n.peak_bytes, mpb)
if mrb > 0:
self.assertGreaterEqual(n.residual_bytes, mrb)
if mob > 0:
self.assertGreaterEqual(n.output_bytes, mob)
check_min(n.children, mm, mam, mcm, mb, mpb, mrb, mob)
with session.Session(config=self._no_rewrite_session_config()) as sess:
x = lib.BuildSmallModel()
self.evaluate(variables.global_variables_initializer())
run_meta = config_pb2.RunMetadata()
_ = sess.run(
x,
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE),
run_metadata=run_meta)
min_val = random.randint(0, 10000)
opts = builder(builder.time_and_memory(
min_micros=min_val)).with_empty_output().build()
tfprof_node = model_analyzer.profile(
sess.graph, run_meta=run_meta, options=opts)
check_min(tfprof_node.children, mm=min_val)
opts = builder(builder.time_and_memory(
min_accelerator_micros=min_val)).with_empty_output().build()
tfprof_node = model_analyzer.profile(
sess.graph, run_meta=run_meta, options=opts)
check_min(tfprof_node.children, mam=min_val)
opts = builder(builder.time_and_memory(
min_cpu_micros=min_val)).with_empty_output().build()
tfprof_node = model_analyzer.profile(
sess.graph, run_meta=run_meta, options=opts)
check_min(tfprof_node.children, mcm=min_val)
opts = builder(builder.time_and_memory(
min_bytes=min_val)).with_empty_output().build()
tfprof_node = model_analyzer.profile(
sess.graph, run_meta=run_meta, options=opts)
check_min(tfprof_node.children, mb=min_val)
opts = builder(builder.time_and_memory(
min_peak_bytes=min_val)).with_empty_output().build()
tfprof_node = model_analyzer.profile(
sess.graph, run_meta=run_meta, options=opts)
check_min(tfprof_node.children, mpb=min_val)
opts = builder(builder.time_and_memory(
min_residual_bytes=min_val)).with_empty_output().build()
tfprof_node = model_analyzer.profile(
sess.graph, run_meta=run_meta, options=opts)
check_min(tfprof_node.children, mrb=min_val)
opts = builder(builder.time_and_memory(
min_output_bytes=min_val)).with_empty_output().build()
tfprof_node = model_analyzer.profile(
sess.graph, run_meta=run_meta, options=opts)
check_min(tfprof_node.children, mob=min_val)
def testSelectOption(self):
ops.reset_default_graph()
outfile = os.path.join(test.get_temp_dir(), 'dump')
def check_selection(selected, not_selected):
with gfile.Open(outfile, 'r') as f:
s = f.read()
for attr in selected:
self.assertTrue(s.find(attr) > 0, s)
for attr in not_selected:
self.assertFalse(s.find(attr) > 0, s)
with session.Session(config=self._no_rewrite_session_config()) as sess:
x = lib.BuildSmallModel()
self.evaluate(variables.global_variables_initializer())
run_meta = config_pb2.RunMetadata()
_ = sess.run(
x,
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE),
run_metadata=run_meta)
opts = builder(
builder.time_and_memory()).with_file_output(outfile).select(
['micros']).build()
_ = model_analyzer.profile(sess.graph, run_meta=run_meta, options=opts)
check_selection(['total execution time', 'accelerator execution time'],
['bytes'])
opts = builder(
builder.time_and_memory()).with_file_output(outfile).select(
['bytes']).build()
_ = model_analyzer.profile(sess.graph, run_meta=run_meta, options=opts)
check_selection(['requested bytes'],
['peak bytes', 'residual bytes', 'output bytes'])
opts = builder(
builder.time_and_memory()).with_file_output(outfile).select(
['peak_bytes', 'residual_bytes', 'output_bytes']).build()
_ = model_analyzer.profile(sess.graph, run_meta=run_meta, options=opts)
check_selection(['peak bytes', 'residual bytes', 'output bytes'],
['requested_bytes'])
def _trainLoop(self, train_op, train_steps, time_dir, time_step, memory_dir,
memory_step, profile_dir, dump_step):
with session.Session(config=self._no_rewrite_session_config()) as sess:
self.evaluate(variables.global_variables_initializer())
# start from 1 because variable_initializer took one step.
for i in range(1, train_steps + 1):
_ = self.evaluate(train_op)
if i in time_step:
ret = gfile.ListDirectory(time_dir)
self.assertEqual(len(ret), 1)
self.assertTrue(
gfile.Open(os.path.join(time_dir, ret[0]), 'r').read().find(
'execution time') > 0)
_ = [gfile.Remove(os.path.join(time_dir, x)) for x in ret]
else:
self.assertEqual(len(gfile.ListDirectory(time_dir)), 0)
if i in memory_step:
ret = gfile.ListDirectory(memory_dir)
self.assertEqual(len(ret), 1)
self.assertTrue(
gfile.Open(os.path.join(memory_dir, ret[0]), 'r').read().find(
'requested bytes') > 0)
_ = [gfile.Remove(os.path.join(memory_dir, x)) for x in ret]
else:
self.assertEqual(len(gfile.ListDirectory(memory_dir)), 0)
if i in dump_step:
ret = gfile.ListDirectory(profile_dir)
self.assertAllEqual(ret, ['profile_%d' % i])
_ = [gfile.Remove(os.path.join(profile_dir, x)) for x in ret]
else:
if i < dump_step[0]:
self.assertFalse(gfile.Exists(profile_dir))
else:
self.assertEqual(len(gfile.ListDirectory(profile_dir)), 0)
@test_util.run_v1_only('b/120545219')
def testAutoProfiling(self):
ops.reset_default_graph()
time_dir = os.path.join(test.get_temp_dir(), 'time')
memory_dir = os.path.join(test.get_temp_dir(), 'memory')
profile_dir = os.path.join(test.get_temp_dir(), 'dir/dir2/profile')
# TODO(xpan): Should we create parent directory for them?
gfile.MkDir(time_dir)
gfile.MkDir(memory_dir)
time_opts = (
builder(builder.time_and_memory()).with_file_output(
os.path.join(time_dir, 'profile')).select(['micros']).build())
memory_opts = (
builder(builder.time_and_memory()).with_file_output(
os.path.join(memory_dir, 'profile')).select(['bytes']).build())
time_steps = [2, 3]
memory_steps = [1, 3]
dump_steps = [3, 4]
x = lib.BuildSmallModel()
with profile_context.ProfileContext(
profile_dir, trace_steps=[1, 2, 3], dump_steps=[3, 4]) as pctx:
pctx.add_auto_profiling('scope', time_opts, time_steps)
pctx.add_auto_profiling('scope', memory_opts, memory_steps)
self._trainLoop(x, 10, time_dir, time_steps, memory_dir, memory_steps,
profile_dir, dump_steps)
@test_util.run_v1_only('b/120545219')
def testOOM(self):
if not test.is_gpu_available():
return
ops.reset_default_graph()
with ops.device('/device:GPU:0'):
a = random_ops.random_normal([1, 10000, 20000], name='test_random1')
b = random_ops.random_normal([30000, 10000, 1], name='test_random2')
c = a * b
try:
with session.Session(config=self._no_rewrite_session_config()) as sess:
sess.run(
c,
options=config_pb2.RunOptions(
report_tensor_allocations_upon_oom=True))
except Exception as e: # pylint: disable=broad-except
exception_str = '%s' % e
# This trace reports allocations for to random tensor.
self.assertTrue('OOM when allocating tensor with shape[30000,10000,20000]'
in exception_str)
mat = re.search('(.*)GiB from test_random2/RandomStandardNormal',
exception_str)
self.assertGreater(float(mat.group(1)), 0.0)
mat = re.search('(.*)MiB from test_random1/RandomStandardNormal',
exception_str)
self.assertGreater(float(mat.group(1)), 0.0)
@test_util.run_v1_only('b/120545219')
def testDistributedOOM(self):
if not test.is_gpu_available():
return
ops.reset_default_graph()
workers, _ = test_util.create_local_cluster(2, 0)
with ops.device('/job:worker/replica:0/task:0/gpu:0'):
a = random_ops.random_normal([1, 10000, 20000], name='test_random1')
with ops.device('/job:worker/replica:0/task:1/gpu:0'):
b = random_ops.random_normal([30000, 10000, 1], name='test_random2')
c = a * b
try:
with session.Session(workers[1].target) as sess:
sess.run(
c,
options=config_pb2.RunOptions(
report_tensor_allocations_upon_oom=True))
except Exception as e: # pylint: disable=broad-except
exception_str = '%s' % e
# test_random2 is reported because it's allocated in worker 1.
self.assertTrue('Current usage from device: '
'/job:worker/replica:0/task:1/device:GPU:0, '
'allocator: GPU_0_bfc' in exception_str)
mat = re.search('(.*)GiB from test_random2/RandomStandardNormal',
exception_str)
self.assertGreater(float(mat.group(1)), 0.0)
# test_random1 is not reported because it's allocated in worker 0.
mat = re.search('(.*)MiB from test_random1/RandomStandardNormal',
exception_str)
self.assertTrue(mat is None)
@test_util.run_v1_only('b/120545219')
def testTrackPersistentBytes(self):
ops.reset_default_graph()
a = array_ops.constant(np.ones((100, 100)))
b = array_ops.constant(np.ones((100, 100)))
c = a * b
config = config_pb2.ConfigProto()
config.graph_options.rewrite_options.min_graph_nodes = -1
with session.Session(config=config) as sess:
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
run_metadata = config_pb2.RunMetadata()
sess.run(c, options=run_options, run_metadata=run_metadata)
options = option_builder.ProfileOptionBuilder.time_and_memory()
options['min_bytes'] = 0
options['select'] = ('bytes', 'peak_bytes', 'output_bytes',
'residual_bytes')
ret = model_analyzer.profile(
sess.graph, run_meta=run_metadata, cmd='scope', options=options)
run_metadata = config_pb2.RunMetadata()
sess.run(c, options=run_options, run_metadata=run_metadata)
ret2 = model_analyzer.profile(
sess.graph, run_meta=run_metadata, cmd='scope', options=options)
n = lib.SearchTFProfNode(ret, 'mul')
n2 = lib.SearchTFProfNode(ret2, 'mul')
self.assertGreater(n.peak_bytes, 0)
self.assertGreater(n.output_bytes, 0)
self.assertGreater(n.residual_bytes, 0)
self.assertEqual(n.peak_bytes, n2.peak_bytes)
self.assertEqual(n.output_bytes, n2.output_bytes)
self.assertEqual(n.residual_bytes, n2.residual_bytes)
@test_util.run_v1_only('b/120545219')
def testTraceLoopBytes(self):
if not test.is_gpu_available():
return
ops.reset_default_graph()
steps = 100
with ops.device('/gpu:0'):
x = array_ops.ones((100, 100), dtype=dtypes.float32)
n = array_ops.constant(steps, dtype=dtypes.int32)
x1 = array_ops.ones((100, 100))
x *= x1
def loop_body(i, x):
x *= x
return i + 1, x
_, y = while_loop.while_loop(lambda i, x: i < n, loop_body,
[array_ops.constant(0), x])
grad = gradients.gradients(y, [x1])
with session.Session(config=self._no_rewrite_session_config()) as sess:
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
run_metadata = config_pb2.RunMetadata()
sess.run(grad, options=run_options, run_metadata=run_metadata)
options = option_builder.ProfileOptionBuilder.time_and_memory()
options['min_bytes'] = 0
options['min_micros'] = 0
options['select'] = ('bytes', 'peak_bytes', 'output_bytes',
'residual_bytes')
options['output'] = 'none'
ret_pb = model_analyzer.profile(
sess.graph, run_meta=run_metadata, cmd='scope', options=options)
self.assertGreater(ret_pb.total_requested_bytes, 1000000)
if __name__ == '__main__':
test.main()
@@ -0,0 +1,461 @@
# Copyright 2016 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.
# ==============================================================================
"""Utilities for building profiler options."""
import copy
from tensorflow.python.profiler import tfprof_logger
from tensorflow.python.util.tf_export import tf_export
@tf_export(v1=['profiler.ProfileOptionBuilder'])
class ProfileOptionBuilder(object):
# pylint: disable=line-too-long
"""Option Builder for Profiling API.
For tutorial on the options, see
https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core/profiler/g3doc/options.md
```python
# Users can use pre-built options:
opts = (
tf.profiler.ProfileOptionBuilder.trainable_variables_parameter())
# Or, build your own options:
opts = (tf.compat.v1.profiler.ProfileOptionBuilder()
.with_max_depth(10)
.with_min_micros(1000)
.select(['accelerator_micros'])
.with_stdout_output()
.build()
# Or customize the pre-built options:
opts = (tf.compat.v1.profiler.ProfileOptionBuilder(
tf.profiler.ProfileOptionBuilder.time_and_memory())
.with_displaying_options(show_name_regexes=['.*rnn.*'])
.build())
# Finally, profiling with the options:
_ = tf.compat.v1.profiler.profile(tf.compat.v1.get_default_graph(),
run_meta=run_meta,
cmd='scope',
options=opts)
```
"""
# pylint: enable=line-too-long
def __init__(self, options=None):
"""Constructor.
Args:
options: Optional initial option dict to start with.
"""
if options is not None:
self._options = copy.deepcopy(options)
else:
self._options = {'max_depth': 100,
'min_bytes': 0,
'min_micros': 0,
'min_params': 0,
'min_float_ops': 0,
'min_occurrence': 0,
'order_by': 'name',
'account_type_regexes': ['.*'],
'start_name_regexes': ['.*'],
'trim_name_regexes': [],
'show_name_regexes': ['.*'],
'hide_name_regexes': [],
'account_displayed_op_only': False,
'select': ['micros'],
'step': -1,
'output': 'stdout'}
@staticmethod
def trainable_variables_parameter():
"""Options used to profile trainable variable parameters.
Normally used together with 'scope' view.
Returns:
A dict of profiling options.
"""
return {'max_depth': 10000,
'min_bytes': 0,
'min_micros': 0,
'min_params': 0,
'min_float_ops': 0,
'min_occurrence': 0,
'order_by': 'name',
'account_type_regexes': [tfprof_logger.TRAINABLE_VARIABLES],
'start_name_regexes': ['.*'],
'trim_name_regexes': [],
'show_name_regexes': ['.*'],
'hide_name_regexes': [],
'account_displayed_op_only': True,
'select': ['params'],
'step': -1,
'output': 'stdout'}
@staticmethod
def float_operation():
# pylint: disable=line-too-long
"""Options used to profile float operations.
Please see https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core/profiler/g3doc/profile_model_architecture.md
on the caveats of calculating float operations.
Returns:
A dict of profiling options.
"""
# pylint: enable=line-too-long
return {'max_depth': 10000,
'min_bytes': 0,
'min_micros': 0,
'min_params': 0,
'min_float_ops': 1,
'min_occurrence': 0,
'order_by': 'float_ops',
'account_type_regexes': ['.*'],
'start_name_regexes': ['.*'],
'trim_name_regexes': [],
'show_name_regexes': ['.*'],
'hide_name_regexes': [],
'account_displayed_op_only': True,
'select': ['float_ops'],
'step': -1,
'output': 'stdout'}
@staticmethod
def time_and_memory(min_micros=1, min_bytes=1, min_accelerator_micros=0,
min_cpu_micros=0, min_peak_bytes=0, min_residual_bytes=0,
min_output_bytes=0):
"""Show operation time and memory consumptions.
Args:
min_micros: Only show profiler nodes with execution time
no less than this. It sums accelerator and cpu times.
min_bytes: Only show profiler nodes requested to allocate no less bytes
than this.
min_accelerator_micros: Only show profiler nodes spend no less than
this time on accelerator (e.g. GPU).
min_cpu_micros: Only show profiler nodes spend no less than
this time on cpu.
min_peak_bytes: Only show profiler nodes using no less than this bytes
at peak (high watermark). For profiler nodes consist of multiple
graph nodes, it sums the graph nodes' peak_bytes.
min_residual_bytes: Only show profiler nodes have no less than
this bytes not being de-allocated after Compute() ends. For
profiler nodes consist of multiple graph nodes, it sums the
graph nodes' residual_bytes.
min_output_bytes: Only show profiler nodes have no less than this bytes
output. The output are not necessarily allocated by this profiler
nodes.
Returns:
A dict of profiling options.
"""
return {'max_depth': 10000,
'min_bytes': min_bytes,
'min_peak_bytes': min_peak_bytes,
'min_residual_bytes': min_residual_bytes,
'min_output_bytes': min_output_bytes,
'min_micros': min_micros,
'min_accelerator_micros': min_accelerator_micros,
'min_cpu_micros': min_cpu_micros,
'min_params': 0,
'min_float_ops': 0,
'min_occurrence': 0,
'order_by': 'micros',
'account_type_regexes': ['.*'],
'start_name_regexes': ['.*'],
'trim_name_regexes': [],
'show_name_regexes': ['.*'],
'hide_name_regexes': [],
'account_displayed_op_only': True,
'select': ['micros', 'bytes'],
'step': -1,
'output': 'stdout'}
def build(self):
"""Build a profiling option.
Returns:
A dict of profiling options.
"""
return copy.deepcopy(self._options)
def with_max_depth(self, max_depth):
"""Set the maximum depth of display.
The depth depends on profiling view. For 'scope' view, it's the
depth of name scope hierarchy (tree), for 'op' view, it's the number
of operation types (list), etc.
Args:
max_depth: Maximum depth of the data structure to display.
Returns:
self
"""
self._options['max_depth'] = max_depth
return self
def with_min_memory(self,
min_bytes=0,
min_peak_bytes=0,
min_residual_bytes=0,
min_output_bytes=0):
"""Only show profiler nodes consuming no less than 'min_bytes'.
Args:
min_bytes: Only show profiler nodes requested to allocate no less bytes
than this.
min_peak_bytes: Only show profiler nodes using no less than this bytes
at peak (high watermark). For profiler nodes consist of multiple
graph nodes, it sums the graph nodes' peak_bytes.
min_residual_bytes: Only show profiler nodes have no less than
this bytes not being de-allocated after Compute() ends. For
profiler nodes consist of multiple graph nodes, it sums the
graph nodes' residual_bytes.
min_output_bytes: Only show profiler nodes have no less than this bytes
output. The output are not necessarily allocated by this profiler
nodes.
Returns:
self
"""
self._options['min_bytes'] = min_bytes
self._options['min_peak_bytes'] = min_peak_bytes
self._options['min_residual_bytes'] = min_residual_bytes
self._options['min_output_bytes'] = min_output_bytes
return self
def with_min_execution_time(self,
min_micros=0,
min_accelerator_micros=0,
min_cpu_micros=0):
"""Only show profiler nodes consuming no less than 'min_micros'.
Args:
min_micros: Only show profiler nodes with execution time
no less than this. It sums accelerator and cpu times.
min_accelerator_micros: Only show profiler nodes spend no less than
this time on accelerator (e.g. GPU).
min_cpu_micros: Only show profiler nodes spend no less than
this time on cpu.
Returns:
self
"""
self._options['min_micros'] = min_micros
self._options['min_accelerator_micros'] = min_accelerator_micros
self._options['min_cpu_micros'] = min_cpu_micros
return self
def with_min_parameters(self, min_params):
"""Only show profiler nodes holding no less than 'min_params' parameters.
'Parameters' normally refers the weights of in TensorFlow variables.
It reflects the 'capacity' of models.
Args:
min_params: Only show profiler nodes holding number parameters
no less than this.
Returns:
self
"""
self._options['min_params'] = min_params
return self
def with_min_occurrence(self, min_occurrence):
# pylint: disable=line-too-long
"""Only show profiler nodes including no less than 'min_occurrence' graph nodes.
A "node" means a profiler output node, which can be a python line
(code view), an operation type (op view), or a graph node
(graph/scope view). A python line includes all graph nodes created by that
line, while an operation type includes all graph nodes of that type.
Args:
min_occurrence: Only show nodes including no less than this.
Returns:
self
"""
# pylint: enable=line-too-long
self._options['min_occurrence'] = min_occurrence
return self
def with_min_float_operations(self, min_float_ops):
# pylint: disable=line-too-long
"""Only show profiler nodes consuming no less than 'min_float_ops'.
Please see https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core/profiler/g3doc/profile_model_architecture.md
on the caveats of calculating float operations.
Args:
min_float_ops: Only show profiler nodes with float operations
no less than this.
Returns:
self
"""
# pylint: enable=line-too-long
self._options['min_float_ops'] = min_float_ops
return self
def with_accounted_types(self, account_type_regexes):
"""Selectively counting statistics based on node types.
Here, 'types' means the profiler nodes' properties. Profiler by default
consider device name (e.g. /job:xx/.../device:GPU:0) and operation type
(e.g. MatMul) as profiler nodes' properties. User can also associate
customized 'types' to profiler nodes through OpLogProto proto.
For example, user can select profiler nodes placed on gpu:0 with:
`account_type_regexes=['.*gpu:0.*']`
If none of a node's properties match the specified regexes, the node is
not displayed nor accounted.
Args:
account_type_regexes: A list of regexes specifying the types.
Returns:
self.
"""
self._options['account_type_regexes'] = copy.copy(account_type_regexes)
return self
def with_node_names(self,
start_name_regexes=None,
show_name_regexes=None,
hide_name_regexes=None,
trim_name_regexes=None):
"""Regular expressions used to select profiler nodes to display.
After 'with_accounted_types' is evaluated, 'with_node_names' are
evaluated as follows:
For a profile data structure, profiler first finds the profiler
nodes matching 'start_name_regexes', and starts displaying profiler
nodes from there. Then, if a node matches 'show_name_regexes' and
doesn't match 'hide_name_regexes', it's displayed. If a node matches
'trim_name_regexes', profiler stops further searching that branch.
Args:
start_name_regexes: list of node name regexes to start displaying.
show_name_regexes: list of node names regexes to display.
hide_name_regexes: list of node_names regexes that should be hidden.
trim_name_regexes: list of node name regexes from where to stop.
Returns:
self
"""
if start_name_regexes is not None:
self._options['start_name_regexes'] = copy.copy(start_name_regexes)
if show_name_regexes is not None:
self._options['show_name_regexes'] = copy.copy(show_name_regexes)
if hide_name_regexes is not None:
self._options['hide_name_regexes'] = copy.copy(hide_name_regexes)
if trim_name_regexes is not None:
self._options['trim_name_regexes'] = copy.copy(trim_name_regexes)
return self
def account_displayed_op_only(self, is_true):
"""Whether only account the statistics of displayed profiler nodes.
Args:
is_true: If true, only account statistics of nodes eventually
displayed by the outputs.
Otherwise, a node's statistics are accounted by its parents
as long as it's types match 'account_type_regexes', even if
it is hidden from the output, say, by hide_name_regexes.
Returns:
self
"""
self._options['account_displayed_op_only'] = is_true
return self
def with_empty_output(self):
"""Do not generate side-effect outputs."""
self._options['output'] = 'none'
return self
def with_stdout_output(self):
"""Print the result to stdout."""
self._options['output'] = 'stdout'
return self
def with_file_output(self, outfile):
"""Print the result to a file."""
self._options['output'] = 'file:outfile=%s' % outfile
return self
def with_timeline_output(self, timeline_file):
"""Generate a timeline json file."""
self._options['output'] = 'timeline:outfile=%s' % timeline_file
return self
def with_pprof_output(self, pprof_file):
"""Generate a pprof profile gzip file.
To use the pprof file:
pprof -png --nodecount=100 --sample_index=1 <pprof_file>
Args:
pprof_file: filename for output, usually suffixed with .pb.gz.
Returns:
self.
"""
self._options['output'] = 'pprof:outfile=%s' % pprof_file
return self
def order_by(self, attribute):
# pylint: disable=line-too-long
"""Order the displayed profiler nodes based on a attribute.
Supported attribute includes micros, bytes, occurrence, params, etc.
https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core/profiler/g3doc/options.md
Args:
attribute: An attribute the profiler node has.
Returns:
self
"""
# pylint: enable=line-too-long
self._options['order_by'] = attribute
return self
def select(self, attributes):
# pylint: disable=line-too-long
"""Select the attributes to display.
See https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core/profiler/g3doc/options.md
for supported attributes.
Args:
attributes: A list of attribute the profiler node has.
Returns:
self
"""
# pylint: enable=line-too-long
self._options['select'] = copy.copy(attributes)
return self
def with_step(self, step):
"""Which profile step to use for profiling.
The 'step' here refers to the step defined by `Profiler.add_step()` API.
Args:
step: When multiple steps of profiles are available, select which step's
profile to use. If -1, use average of all available steps.
Returns:
self
"""
self._options['step'] = step
return self
@@ -0,0 +1,440 @@
# 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.
# ==============================================================================
"""Profiler for TensorFlow models that outputs data in pprof format.
See https://github.com/google/pprof/blob/master/proto/profile.proto for pprof
profile format.
The following needs to be set for profiler to work:
* trace_level needs to be set to FULL_TRACE
* run_metadata object should be passed in to session.run call
Sample usage:
options = tf.compat.v1.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.compat.v1.RunMetadata()
with tf.compat.v1.Session as sess:
...
sess.run(computation, run_metadata=run_metadata, options=options)
pprof_profiler.profile(sess.graph, run_metadata, output_dir)
The code above would output a pprof profile to separate output_dir/.*.pb.gz
file for each device. These files can be passed to pprof for formatting.
For e.g.:
pprof -png --nodecount=100 --sample_index=1 output_dir/profile_output.pb.gz
"""
from collections import defaultdict
from collections import namedtuple
import gzip
import os
import string
import sys
import time
from proto import profile_pb2
if sys.version_info < (3,):
maketrans = string.maketrans
else:
maketrans = str.maketrans
ProfileDatum = namedtuple('ProfileDatum', [
'node_exec_stats', 'op_type', 'traceback'])
class StringTable(object):
"""Keeps track of strings to add to string_table in pprof proto."""
def __init__(self):
# Pprof requires first entry in string_table to be ''.
self._string_table = ['']
self._string_to_index = {'': 0}
def index_of(self, value_str):
"""Get index of value_str in the string table.
If value_str is not in the string table, we will add it at the end
and then return the new index.
Args:
value_str: (string) Value to lookup/add in/to the string table.
Returns:
Index of value_str in the string table.
"""
if value_str is None:
value_str = ''
if value_str in self._string_to_index:
return self._string_to_index[value_str]
index = len(self._string_table)
self._string_table.append(value_str)
self._string_to_index[value_str] = index
return index
def next_index(self):
"""Gets index that would be assigned to the next added string.
Returns:
Index of the next string if it was added.
"""
return len(self._string_table)
def string_table(self):
"""Returns a list of strings to store in pprof's string_table."""
return self._string_table
class Functions(object):
"""Keeps track of `Function` protos for pprof profile."""
def __init__(self, string_table):
"""Constructor.
Args:
string_table: A `StringTable` object.
"""
self._string_table = string_table
# Maps tuples in the form (file_path, function_name, start_line_number)
# to `Function` protos.
self._function_key_to_function = {}
def index_of(self, file_path, function_name, function_start_line):
"""Returns index of the function, adding the function if needed.
Args:
file_path: (string) Path to file where the function is defined.
function_name: (string) Function name.
function_start_line: (integer) Start line number of function definition.
Returns:
Function index.
"""
function_key = (file_path, function_name, function_start_line)
if function_key in self._function_key_to_function:
return self._function_key_to_function[function_key].id
else:
# Function indexes should start from 1
function_index = len(self._function_key_to_function) + 1
function = profile_pb2.Function()
function.id = function_index
function.name = self._string_table.index_of(function_name)
function.filename = self._string_table.index_of(file_path)
function.start_line = function_start_line
self._function_key_to_function[function_key] = function
return function_index
def function_protos(self):
"""Returns list of `profile_pb2.Function` protos."""
return self._function_key_to_function.values()
class Locations(object):
"""Keeps track of `Location` protos for pprof profile.
`Locations` store information about function call locations.
"""
def __init__(self, functions):
"""Constructor.
Args:
functions: A `Functions` object.
"""
self._functions = functions
# Maps tuples in the form (file_path, called_function_name, line_number)
# to `Location` protos.
self._location_key_to_location = {}
def index_of(
self, file_path, line_number, called_function_name, called_file_path,
called_function_start_line):
"""Returns index of the location, adding the location if needed.
Args:
file_path: (string) Path to file that makes the call.
line_number: (integer) Call line number.
called_function_name: (string) Function name of the function called at
`file_path` and `line_number`.
called_file_path: (string) Path to file where the called function is
defined.
called_function_start_line: (integer) Start line number of called
function definition in `called_file_path` file.
Returns:
Index of location.
"""
location_key = (file_path, called_function_name, line_number)
if location_key in self._location_key_to_location:
location = self._location_key_to_location[location_key]
return location.id
else:
# Location indexes should start from 1
location_index = len(self._location_key_to_location) + 1
location = profile_pb2.Location()
location.id = location_index
self._location_key_to_location[location_key] = location
line = location.line.add()
line.function_id = self._functions.index_of(
called_file_path, called_function_name, called_function_start_line)
line.line = line_number
return location_index
def location_protos(self):
"""Returns list of `profile_pb2.Location` protos."""
return self._location_key_to_location.values()
class Samples(object):
"""Keeps track of `Sample` protos for pprof profile.
Samples store the following statistics in order:
count, all_time, op_time
"""
def __init__(self, string_table):
"""Constructor.
Args:
string_table: A `StringTable` object.
"""
self._string_table = string_table
# TODO(annarev): figure out if location is unique for each node name.
# If not, also key this dictionary based on location ids.
self._node_name_to_sample = {}
def add(self, datum, location_ids):
"""Adds a sample data point.
Args:
datum: `ProfileDatum` to add a sample for.
location_ids: List of numeric location ids for this sample.
"""
node_name = datum.node_exec_stats.node_name
if node_name in self._node_name_to_sample:
sample = self._node_name_to_sample[node_name]
sample.location_id.extend(location_ids)
else:
sample = profile_pb2.Sample()
# Sample stores 3 values: count, all_time, op_time
sample.value.extend([0, 0, 0])
label = sample.label.add()
label.key = self._string_table.index_of('node_name')
label.str = self._string_table.index_of(node_name)
label = sample.label.add()
label.key = self._string_table.index_of('op_type')
label.str = self._string_table.index_of(datum.op_type)
self._node_name_to_sample[node_name] = sample
sample.value[0] += 1
sample.value[1] += datum.node_exec_stats.all_end_rel_micros
sample.value[2] += (
datum.node_exec_stats.op_end_rel_micros -
datum.node_exec_stats.op_start_rel_micros)
def get_sample_protos(self):
"""Returns list of `Sample` protos for pprof profile."""
return self._node_name_to_sample.values()
class PprofProfiler(object):
"""Creates profiles in pprof format."""
def __init__(self, graph, run_metadata):
"""Constructor.
Args:
graph: A `Graph` instance.
run_metadata: A list of `RunMetadata` objects.
"""
self._graph = graph
self._run_metadata = run_metadata
self._string_table = StringTable()
self._functions = Functions(self._string_table)
self._locations = Locations(self._functions)
def profile(self):
"""Generates pprof profiles.
Returns:
Dictionary mapping from device name to proto in `profile_pb2.Profile`
format.
"""
profiles = {}
data_generator_func = self._get_profile_data_generator()
for device_index, device_stats in enumerate(
self._run_metadata.step_stats.dev_stats):
# Create profile
pprof_proto = self._get_pprof_proto(data_generator_func(device_stats))
if not pprof_proto.sample:
print(
'Not enough data to create profile for device %s. Did you pass '
'RunMetadata to session.run call?' % device_stats.device)
continue
# Add device name comment
device_count = len(self._run_metadata.step_stats.dev_stats)
device_description = (
'Device %d of %d: %s' %
(device_index + 1, device_count, device_stats.device))
device_description_str_index = self._string_table.next_index()
pprof_proto.string_table.append(device_description)
pprof_proto.comment.append(device_description_str_index)
profiles[device_stats.device] = pprof_proto
return profiles
def _get_pprof_proto(self, profile_datum_generator):
"""Returns profile data in pprof proto format.
Args:
profile_datum_generator: Generator outputting `ProfileDatum` objects.
Returns:
A proto in pprof format.
"""
pprof_profile = profile_pb2.Profile()
samples = Samples(self._string_table)
for datum in profile_datum_generator:
if not datum.traceback:
continue
stack_frame = datum.traceback[-1]
after_apply_op = False
location_ids = []
# We add locations from stack trace in bottom-up order.
for stack_frame_index in reversed(range(len(datum.traceback) - 1)):
prev_stack_frame = stack_frame
stack_frame = datum.traceback[stack_frame_index]
# Call at current frame calls function at previous frame.
prev_file_path = prev_stack_frame[0]
prev_function = prev_stack_frame[2]
prev_function_start_line = -1
curr_file_path = stack_frame[0]
curr_line_number = stack_frame[1]
# Skip all calls up to apply_op since they are the same for all ops.
if not after_apply_op:
if prev_function == 'apply_op':
after_apply_op = True
continue
location_index = self._locations.index_of(
curr_file_path, curr_line_number,
prev_function, prev_file_path, prev_function_start_line)
location_ids.append(location_index)
samples.add(datum, location_ids)
sample_type_description = 'count'
sample_type = pprof_profile.sample_type.add()
sample_type.type = self._string_table.index_of(sample_type_description)
sample_type.unit = self._string_table.index_of('count')
sample_type_description = 'all_time'
sample_type = pprof_profile.sample_type.add()
sample_type.type = self._string_table.index_of(sample_type_description)
sample_type.unit = self._string_table.index_of('nanoseconds')
sample_type_description = 'op_time'
sample_type = pprof_profile.sample_type.add()
sample_type.type = self._string_table.index_of(sample_type_description)
sample_type.unit = self._string_table.index_of('nanoseconds')
pprof_profile.string_table.extend(self._string_table.string_table())
pprof_profile.sample.extend(samples.get_sample_protos())
pprof_profile.function.extend(self._functions.function_protos())
pprof_profile.location.extend(self._locations.location_protos())
return pprof_profile
def _get_profile_data_generator(self):
"""Get function that generates `ProfileDatum` objects.
Returns:
A function that generates `ProfileDatum` objects.
"""
node_to_traceback = defaultdict(list)
node_to_op_type = defaultdict(str)
for op in self._graph.get_operations():
node_to_traceback[op.name] = op.traceback
node_to_op_type[op.name] = op.type
def profile_data_generator(device_step_stats):
for node_stats in device_step_stats.node_stats:
if node_stats.node_name == '_SOURCE' or node_stats.node_name == '_SINK':
continue
yield ProfileDatum(
node_stats,
node_to_op_type[node_stats.node_name],
node_to_traceback[node_stats.node_name])
return profile_data_generator
def get_profiles(graph, run_metadata):
"""Generate profiles in pprof format.
See https://github.com/google/pprof/blob/master/proto/profile.proto
for pprof proto format.
Args:
graph: A `Graph` object.
run_metadata: A `RunMetadata` proto.
Returns:
A dictionary mapping from device name to pprof proto for that device.
"""
return PprofProfiler(graph, run_metadata).profile()
def profile(graph, run_metadata, output_dir=None):
"""Generate profiles in pprof format.
See https://github.com/google/pprof/blob/master/proto/profile.proto
for pprof proto format.
Args:
graph: A `Graph` object.
run_metadata: A `RunMetadata` proto.
output_dir: (string) Directory to output pprof profile to.
Profile files for each device will be stored in compressed
serialized proto format. If output_dir is None, profile protos
will be printed to stdout instead.
Returns:
List of output files created by this profile call.
(Note: this list will be empty if output_dir is None)
"""
profiles = get_profiles(graph, run_metadata)
output_file_template = None
if output_dir:
if not os.path.isdir(output_dir):
os.makedirs(output_dir)
time_suffix = time.strftime('%Y%m%d%H%M%S')
output_file_template = os.path.join(
output_dir, '%s_' + time_suffix + '.pb.gz')
profile_files = []
for device, pprof_proto in profiles.items():
if output_file_template is None:
print('No output directory specified, printing to stdout instead.')
print(pprof_proto)
else:
device_name = str(device).strip('/').translate(
maketrans('/:', '__'))
profile_file = output_file_template % device_name
profile_files.append(profile_file)
with gzip.open(profile_file, 'w') as output_file:
print('Writing profile to %s...' % profile_file)
output_file.write(pprof_proto.SerializeToString())
return profile_files
@@ -0,0 +1,162 @@
# 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.
# ==============================================================================
"""Tests for pprof_profiler."""
import gzip
from proto import profile_pb2
from tensorflow.core.framework import step_stats_pb2
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import test_util
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import while_loop
from tensorflow.python.platform import test
from tensorflow.python.profiler import pprof_profiler
class PprofProfilerTest(test.TestCase):
def testDataEmpty(self):
output_dir = test.get_temp_dir()
run_metadata = config_pb2.RunMetadata()
graph = test.mock.MagicMock()
graph.get_operations.return_value = []
profiles = pprof_profiler.get_profiles(graph, run_metadata)
self.assertEqual(0, len(profiles))
profile_files = pprof_profiler.profile(
graph, run_metadata, output_dir)
self.assertEqual(0, len(profile_files))
def testRunMetadataEmpty(self):
output_dir = test.get_temp_dir()
run_metadata = config_pb2.RunMetadata()
graph = test.mock.MagicMock()
op1 = test.mock.MagicMock()
op1.name = 'Add/123'
op1.traceback = [('a/b/file1', 10, 'some_var')]
op1.type = 'add'
graph.get_operations.return_value = [op1]
profiles = pprof_profiler.get_profiles(graph, run_metadata)
self.assertEqual(0, len(profiles))
profile_files = pprof_profiler.profile(
graph, run_metadata, output_dir)
self.assertEqual(0, len(profile_files))
def testValidProfile(self):
output_dir = test.get_temp_dir()
run_metadata = config_pb2.RunMetadata()
node1 = step_stats_pb2.NodeExecStats(
node_name='Add/123',
op_start_rel_micros=3,
op_end_rel_micros=5,
all_end_rel_micros=4)
run_metadata = config_pb2.RunMetadata()
device1 = run_metadata.step_stats.dev_stats.add()
device1.device = 'deviceA'
device1.node_stats.extend([node1])
graph = test.mock.MagicMock()
op1 = test.mock.MagicMock()
op1.name = 'Add/123'
op1.traceback = [
('a/b/file1', 10, 'apply_op', 'abc'), ('a/c/file2', 12, 'my_op', 'def')]
op1.type = 'add'
graph.get_operations.return_value = [op1]
expected_proto = """sample_type {
type: 5
unit: 5
}
sample_type {
type: 6
unit: 7
}
sample_type {
type: 8
unit: 7
}
sample {
value: 1
value: 4
value: 2
label {
key: 1
str: 2
}
label {
key: 3
str: 4
}
}
string_table: ""
string_table: "node_name"
string_table: "Add/123"
string_table: "op_type"
string_table: "add"
string_table: "count"
string_table: "all_time"
string_table: "nanoseconds"
string_table: "op_time"
string_table: "Device 1 of 1: deviceA"
comment: 9
"""
# Test with protos
profiles = pprof_profiler.get_profiles(graph, run_metadata)
self.assertEqual(1, len(profiles))
self.assertTrue('deviceA' in profiles)
self.assertEqual(expected_proto, str(profiles['deviceA']))
# Test with files
profile_files = pprof_profiler.profile(
graph, run_metadata, output_dir)
self.assertEqual(1, len(profile_files))
with gzip.open(profile_files[0]) as profile_file:
profile_contents = profile_file.read()
profile = profile_pb2.Profile()
profile.ParseFromString(profile_contents)
self.assertEqual(expected_proto, str(profile))
@test_util.run_v1_only('b/120545219')
def testProfileWithWhileLoop(self):
options = config_pb2.RunOptions()
options.trace_level = config_pb2.RunOptions.FULL_TRACE
run_metadata = config_pb2.RunMetadata()
num_iters = 5
with self.cached_session() as sess:
i = constant_op.constant(0)
c = lambda i: math_ops.less(i, num_iters)
b = lambda i: math_ops.add(i, 1)
r = while_loop.while_loop(c, b, [i])
sess.run(r, options=options, run_metadata=run_metadata)
profiles = pprof_profiler.get_profiles(sess.graph, run_metadata)
self.assertEqual(1, len(profiles))
profile = next(iter(profiles.values()))
add_samples = [] # Samples for the while/Add node
for sample in profile.sample:
if profile.string_table[sample.label[0].str] == 'while/Add':
add_samples.append(sample)
# Values for same nodes are aggregated.
self.assertEqual(1, len(add_samples))
# Value of "count" should be equal to number of iterations.
self.assertEqual(num_iters, add_samples[0].value[0])
if __name__ == '__main__':
test.main()
@@ -0,0 +1,356 @@
# Copyright 2016 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.
# ==============================================================================
"""A Context that captures profile and performs profiling/dumping.
"""
import contextlib
import os
import random
import sys
import threading
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.client import session
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.platform import gfile
from tensorflow.python.profiler import model_analyzer
from tensorflow.python.util import _pywrap_tfprof as print_mdl
from tensorflow.python.util import compat
WARMUP_STEPS = 10
MAX_TRACED_STEPS = 100
def _profiled_init(self, target='', graph=None, config=None):
"""Overwrites the session.__init__."""
self._profiler_init_internal(target, graph, config) # pylint: disable=protected-access
def _profiled_run(self,
fetches,
feed_dict=None,
options=None,
run_metadata=None):
"""Overwrites the session.run()."""
# pylint: disable=protected-access
# Count the session steps.
with self.profile_context._new_step() as state:
step, locked = state
# Fast path if no need for profiling.
if locked and not self.profile_context._is_fast_path(step):
# Maybe trace this step.
if self.profile_context._should_trace(step, self.graph, fetches):
if self.profile_context._debug:
sys.stderr.write('debug: tracing step: %d\n' % step)
# Enable tracing, perform auto profiling or auto dump.
if not run_metadata:
run_metadata = config_pb2.RunMetadata()
if not options:
options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
old_trace_level = options.trace_level
else:
old_trace_level = options.trace_level
options.trace_level = config_pb2.RunOptions.FULL_TRACE
ret = self._profiler_run_internal(
fetches, feed_dict, options, run_metadata)
if self.profile_context._debug:
self.profile_context._dump_file(run_metadata, 'run_meta_%d' % step)
self.profile_context.profiler._graph = self.graph
self.profile_context.profiler.add_step(step, run_metadata)
options.trace_level = old_trace_level
else:
ret = self._profiler_run_internal(fetches, feed_dict, options)
# Maybe dump profile.
self.profile_context._maybe_dump(step)
# Maybe profile:
to_profiles = self.profile_context._profile_candidates()
for to_prof in to_profiles:
cmd, opts, _ = to_prof
saved_views = self.profile_context._views.setdefault(cmd, {})
if self.profile_context._debug:
sys.stderr.write('debug: profiling %s step: %d\n' % (cmd, step))
if cmd == 'graph':
saved_views[step] = self.profile_context.profiler.profile_graph(opts)
elif cmd == 'scope':
saved_views[step] = self.profile_context.profiler.profile_name_scope(
opts)
elif cmd == 'op':
saved_views[step] = self.profile_context.profiler.profile_operations(
opts)
elif cmd == 'code':
saved_views[step] = self.profile_context.profiler.profile_python(opts)
else:
raise ValueError('Unknown cmd: %s\n' % cmd)
return ret
# Fast no lock path.
return self._profiler_run_internal(
fetches, feed_dict, options, run_metadata)
# pylint: enable=protected-access
class ProfileContext(object):
"""A Context that captures RunMetadata and performs profiling.
```python
# Trace steps 100~200, profile at [150, 200] and dump profile at 200.
with profile_context.ProfileContext('/tmp/train_dir',
trace_steps=range(100, 200, 3),
dump_steps=[200]) as pctx:
opts = tf.profiler.ProfileOptionBuilder.time_and_memory()
pctx.add_auto_profiling('op', opts, [150, 200])
train_loop().
# Tracing only.
with profile_context.tfprof.ProfileContext('/tmp/train_dir') as pctx:
# Run train/eval loop for at least few hundred steps. Profiles will be
# dumped to train_dir. Use web UI or command line to do profiling.
train_loop().
# When session object is available, do explicit trace, profile and dump.
with profile_context.ProfileContext('/tmp/train_dir',
trace_steps=[],
dump_steps=[]) as pctx:
opts = tf.profiler.ProfileOptionBuilder.time_and_memory()
pctx.trace_next_step()
_ = session.run(train_op)
pctx.profiler.profile_operations(options=opts)
```
Args:
profile_dir: Directory to store profiles.
trace_steps: A list of session run steps to trace. If None, use
pre-defined steps.
dump_steps: A list of steps to dump the profile to `profile_dir`. If None,
use pre-defined steps.
enabled: If false, everything is disabled with minimal overhead. It allows
user to only enable profiling when needed.
debug: If true, also dumps the raw trace RunMetadata text file to
profile_dir. And print debugging message. Useful for bug report.
"""
def __init__(self,
profile_dir,
trace_steps=None,
dump_steps=None,
enabled=True,
debug=False):
self._enabled = enabled
if not self._enabled:
return
self._debug = debug
if not profile_dir:
raise ValueError('Must have a directory for profile.\n')
self._profiler_dir = profile_dir
if trace_steps is None:
self._trace_steps = set()
self._auto_tracing = True
else:
if len(trace_steps) > MAX_TRACED_STEPS:
raise ValueError('Only support tracing up to 100 steps.\n')
self._trace_steps = set(trace_steps[:])
self._auto_tracing = False
if dump_steps is None:
self._dump_steps = set([MAX_TRACED_STEPS])
else:
self._dump_steps = set(dump_steps[:])
self._rng = random.Random(111)
self._fetched = set()
self._slow_path_steps = self._dump_steps | self._trace_steps
self._trace_next_step = False
self._dump_next_step = False
self._step = 0
self._traced_steps = 0
self._auto_profiles = []
self._profiler = None
self._views = {}
self._lock = threading.Lock()
def get_profiles(self, cmd):
"""Returns profiling results for each step at which `cmd` was run.
Args:
cmd: string, profiling command used in an `add_auto_profiling` call.
Returns:
dict[int: (MultiGraphNodeProto | GraphNodeProto)]. Keys are steps at which
the profiling command was run. Values are the outputs of profiling.
For "code" and "op" commands this will be a `MultiGraphNodeProto`, for
"scope" and "graph" commands this will be a `GraphNodeProto.
Raises:
ValueError: if `cmd` was never run (either because no session.run call was
made or because there was no `add_auto_profiling` call with the specified
`cmd`.
"""
if cmd not in self._views:
raise ValueError('No autoprofiler for command: {}, was run'.format(cmd))
return self._views[cmd]
def add_auto_profiling(self, cmd, options, profile_steps):
"""Traces and profiles at some session run steps.
Args:
cmd: The profiling commands. (i.e. scope, op, python, graph)
options: The profiling options.
profile_steps: A list/set of integers. The profiling command and options
will be run automatically at these integer steps. Each step is
a session.run.
"""
if not self._enabled:
return
self._auto_profiles.append((cmd, options, profile_steps[:]))
self._slow_path_steps |= set(profile_steps)
self._trace_steps |= set(profile_steps)
@property
def profiler(self):
"""Returns the current profiler object."""
if not self._enabled:
return None
if not self._profiler:
self._profiler = model_analyzer.Profiler(ops.get_default_graph())
return self._profiler
def trace_next_step(self):
"""Enables tracing and adds traces to profiler at next step."""
if not self._enabled:
return
self._trace_next_step = True
self._slow_path_steps.add(self._step)
def dump_next_step(self):
"""Enable tracing and dump profiles at next step."""
if not self._enabled:
return
self._dump_next_step = True
self._slow_path_steps.add(self._step)
def _is_fast_path(self, step):
if step in self._slow_path_steps:
return False
# When user doesn't set the tracing steps explicitly, auto decide it.
if (self._auto_tracing and step > WARMUP_STEPS and
self._traced_steps <= MAX_TRACED_STEPS):
return False
return True
def _should_trace(self, step, graph, fetches):
"""Whether should do tracing at current step."""
if self._traced_steps > MAX_TRACED_STEPS:
return False
# Check user-set tracing steps.
if step in self._trace_steps or self._trace_next_step:
self._traced_steps += 1
return True
# If no user-set tracing steps set and passes warm up steps, auto trace.
if self._auto_tracing and step > WARMUP_STEPS:
# If the fetches have not been seen before, trace it.
with graph.as_default():
fetch_names = [f.name for f in
session._FetchMapper.for_fetch(fetches).unique_fetches()] # pylint: disable=protected-access
fetch_name = '-'.join(sorted(fetch_names))
if self._debug:
sys.stderr.write('debug: trace fetches: %s\n' % fetch_name)
if fetch_name not in self._fetched:
self._fetched.add(fetch_name)
self._traced_steps += 1
return True
# If the trace coverage is low, does some random tracing.
if (self.profiler._coverage < 0.5 and step < MAX_TRACED_STEPS and # pylint: disable=protected-access
self._rng.randint(0, 10) < 2):
self._traced_steps += 1
return True
return False
def _maybe_dump(self, step):
"""Maybe dump the profile file."""
if not (step in self._dump_steps or self._dump_next_step):
return
if self._debug:
sys.stderr.write('debug: dumping file at step: %d\n' % step)
gfile.MakeDirs(self._profiler_dir)
filename = os.path.join(compat.as_bytes(self._profiler_dir),
compat.as_bytes('profile_%d' % step))
self.profiler._write_profile(filename) # pylint: disable=protected-access
def _dump_file(self, pb, basename):
gfile.MakeDirs(self._profiler_dir)
with gfile.Open(os.path.join(self._profiler_dir, basename), 'w') as f:
f.write('%s' % pb)
@contextlib.contextmanager
def _new_step(self):
acquired = self._lock.acquire(False) # pylint: disable=assignment-from-no-return
yield (self._step, acquired)
self._step += 1
self._trace_next_step = False
self._dump_next_step = False
if acquired:
self._lock.release()
def _profile_candidates(self):
to_profile = []
for auto_prof in self._auto_profiles:
_, _, prof_steps = auto_prof
if self._step in prof_steps:
to_profile.append(auto_prof)
return to_profile
def __enter__(self):
if self._enabled:
self.old_run = getattr(session.BaseSession, 'run', None)
self.old_init = getattr(session.BaseSession, '__init__', None)
if not self.old_run:
raise errors.InternalError(None, None, 'BaseSession misses run method.')
elif not self.old_init:
raise errors.InternalError(None, None,
'BaseSession misses __init__ method.')
elif getattr(session.BaseSession, '_profiler_run_internal', None):
raise errors.InternalError(None, None,
'Already in context or context not cleaned.')
elif getattr(session.BaseSession, '_profiler_init_internal', None):
raise errors.InternalError(None, None,
'Already in context or context not cleaned.')
else:
setattr(session.BaseSession, 'run', _profiled_run)
setattr(session.BaseSession, '__init__', _profiled_init)
setattr(session.BaseSession, '_profiler_run_internal', self.old_run)
setattr(session.BaseSession, '_profiler_init_internal', self.old_init)
setattr(session.BaseSession, 'profile_context', self)
return self
else:
return self
def __exit__(self, exec_type, exec_value, exec_tb):
if not self._enabled:
return
print_mdl.DeleteProfiler()
setattr(session.BaseSession, 'run', self.old_run)
setattr(session.BaseSession, '__init__', self.old_init)
setattr(session.BaseSession, '_profiler_run_internal', None)
setattr(session.BaseSession, '_profiler_init_internal', None)
setattr(session.BaseSession, 'profile_context', None)
@@ -0,0 +1,116 @@
# Copyright 2016 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.
# ==============================================================================
import os
from tensorflow.python.client import session
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import variables
from tensorflow.python.platform import gfile
from tensorflow.python.platform import test
from tensorflow.python.profiler import option_builder
# pylint: disable=g-bad-import-order
from tensorflow.python.profiler import profile_context
from tensorflow.python.profiler.internal import model_analyzer_testlib as lib
builder = option_builder.ProfileOptionBuilder
class ProfilerContextTest(test.TestCase):
@test_util.run_deprecated_v1
def testBasics(self):
ops.reset_default_graph()
outfile = os.path.join(test.get_temp_dir(), "dump")
opts = builder(builder.time_and_memory()).with_file_output(outfile).build()
x = lib.BuildFullModel()
profile_str = None
profile_step100 = os.path.join(test.get_temp_dir(), "profile_100")
with profile_context.ProfileContext(test.get_temp_dir()) as pctx:
pctx.add_auto_profiling("op", options=opts, profile_steps=[15, 50, 100])
with session.Session() as sess:
self.evaluate(variables.global_variables_initializer())
total_steps = 101
for i in range(total_steps):
self.evaluate(x)
if i == 14 or i == 49:
self.assertTrue(gfile.Exists(outfile))
gfile.Remove(outfile)
if i == 99:
self.assertTrue(gfile.Exists(profile_step100))
with gfile.Open(outfile, "r") as f:
profile_str = f.read()
gfile.Remove(outfile)
self.assertEqual(set([15, 50, 100]), set(pctx.get_profiles("op").keys()))
with lib.ProfilerFromFile(os.path.join(test.get_temp_dir(),
"profile_100")) as profiler:
profiler.profile_operations(options=opts)
with gfile.Open(outfile, "r") as f:
self.assertEqual(profile_str, f.read())
@test_util.run_deprecated_v1
def testAutoTracingInDeubMode(self):
ops.reset_default_graph()
x = lib.BuildFullModel()
with profile_context.ProfileContext(test.get_temp_dir(), debug=True):
with session.Session() as sess:
self.evaluate(variables.global_variables_initializer())
for _ in range(10):
self.evaluate(x)
for f in gfile.ListDirectory(test.get_temp_dir()):
# Warm up, no tracing.
self.assertFalse("run_meta" in f)
self.evaluate(x)
self.assertTrue(
gfile.Exists(os.path.join(test.get_temp_dir(), "run_meta_11")))
gfile.Remove(os.path.join(test.get_temp_dir(), "run_meta_11"))
# fetched already.
self.evaluate(x)
for f in gfile.ListDirectory(test.get_temp_dir()):
self.assertFalse("run_meta" in f)
@test_util.run_deprecated_v1
def testDisabled(self):
ops.reset_default_graph()
x = lib.BuildFullModel()
with profile_context.ProfileContext(
test.get_temp_dir(), enabled=False) as pctx:
with session.Session() as sess:
self.evaluate(variables.global_variables_initializer())
for _ in range(10):
self.evaluate(x)
self.assertTrue(pctx.profiler is None)
self.assertTrue(
getattr(session.BaseSession, "profile_context", None) is None)
with profile_context.ProfileContext(test.get_temp_dir()) as pctx:
with session.Session() as sess:
self.evaluate(variables.global_variables_initializer())
for _ in range(10):
self.evaluate(x)
self.assertFalse(pctx.profiler is None)
self.assertFalse(
getattr(session.BaseSession, "profile_context", None) is None)
if __name__ == "__main__":
test.main()
+51
View File
@@ -0,0 +1,51 @@
# Copyright 2016 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.
# ==============================================================================
"""profiler python module provides APIs to profile TensorFlow models.
"""
# pylint: disable=unused-import
from tensorflow.core.profiler.tfprof_log_pb2 import OpLogProto
from tensorflow.core.profiler.tfprof_output_pb2 import AdviceProto
from tensorflow.core.profiler.tfprof_output_pb2 import GraphNodeProto
from tensorflow.core.profiler.tfprof_output_pb2 import MultiGraphNodeProto
from tensorflow.python.profiler.model_analyzer import advise
from tensorflow.python.profiler.model_analyzer import profile
from tensorflow.python.profiler.model_analyzer import Profiler
from tensorflow.python.profiler.option_builder import ProfileOptionBuilder
from tensorflow.python.profiler.tfprof_logger import write_op_log
from tensorflow.python.util.tf_export import tf_export
_allowed_symbols = [
'Profiler',
'profile',
'ProfileOptionBuilder',
'advise',
'write_op_log',
]
_allowed_symbols.extend([
'GraphNodeProto',
'MultiGraphNodeProto',
'AdviceProto',
'OpLogProto',
])
# Export protos
tf_export(v1=['profiler.GraphNodeProto'])(GraphNodeProto)
tf_export(v1=['profiler.MultiGraphNodeProto'])(MultiGraphNodeProto)
tf_export(v1=['profiler.AdviceProto'])(AdviceProto)
tf_export(v1=['profiler.OpLogProto'])(OpLogProto)
@@ -0,0 +1,178 @@
# Copyright 2020 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.
# ==============================================================================
"""Profiler client APIs."""
from tensorflow.python.framework import errors
from tensorflow.python.profiler.internal import _pywrap_profiler_plugin
from tensorflow.python.util.tf_export import tf_export
_GRPC_PREFIX = 'grpc://'
@tf_export('profiler.experimental.client.trace', v1=[])
def trace(service_addr,
logdir,
duration_ms,
worker_list='',
num_tracing_attempts=3,
options=None):
"""Sends gRPC requests to one or more profiler servers to perform on-demand profiling.
This method will block the calling thread until it receives responses from all
servers or until deadline expiration. Both single host and multiple host
profiling are supported on CPU, GPU, and TPU.
The profiled results will be saved by each server to the specified TensorBoard
log directory (i.e. the directory you save your model checkpoints). Use the
TensorBoard profile plugin to view the visualization and analysis results.
Args:
service_addr: A comma delimited string of gRPC addresses of the workers to
profile.
e.g. service_addr='grpc://localhost:6009'
service_addr='grpc://10.0.0.2:8466,grpc://10.0.0.3:8466'
service_addr='grpc://localhost:12345,grpc://localhost:23456'
logdir: Path to save profile data to, typically a TensorBoard log directory.
This path must be accessible to both the client and server.
e.g. logdir='gs://your_tb_dir'
duration_ms: Duration of tracing or monitoring in milliseconds. Must be
greater than zero.
worker_list: An optional TPU only configuration. The list of workers to
profile in the current session.
num_tracing_attempts: Optional. Automatically retry N times when no trace
event is collected (default 3).
options: profiler.experimental.ProfilerOptions namedtuple for miscellaneous
profiler options.
Raises:
InvalidArgumentError: For when arguments fail validation checks.
UnavailableError: If no trace event was collected.
Example usage (CPU/GPU):
```python
# Start a profiler server before your model runs.
tf.profiler.experimental.server.start(6009)
# (Model code goes here).
# Send gRPC request to the profiler server to collect a trace of your model.
tf.profiler.experimental.client.trace('grpc://localhost:6009',
'/nfs/tb_log', 2000)
```
Example usage (Multiple GPUs):
```python
# E.g. your worker IP addresses are 10.0.0.2, 10.0.0.3, 10.0.0.4, and you
# would like to schedule start of profiling 1 second from now, for a
# duration of 2 seconds.
options['delay_ms'] = 1000
tf.profiler.experimental.client.trace(
'grpc://10.0.0.2:8466,grpc://10.0.0.3:8466,grpc://10.0.0.4:8466',
'gs://your_tb_dir',
2000,
options=options)
```
Example usage (TPU):
```python
# Send gRPC request to a TPU worker to collect a trace of your model. A
# profiler service has been started in the TPU worker at port 8466.
# E.g. your TPU IP address is 10.0.0.2 and you want to profile for 2 seconds
# .
tf.profiler.experimental.client.trace('grpc://10.0.0.2:8466',
'gs://your_tb_dir', 2000)
```
Example usage (Multiple TPUs):
```python
# Send gRPC request to a TPU pod to collect a trace of your model on
# multiple TPUs. A profiler service has been started in all the TPU workers
# at the port 8466.
# E.g. your TPU IP addresses are 10.0.0.2, 10.0.0.3, 10.0.0.4, and you want
# to profile for 2 seconds.
tf.profiler.experimental.client.trace(
'grpc://10.0.0.2:8466',
'gs://your_tb_dir',
2000,
'10.0.0.2:8466,10.0.0.3:8466,10.0.0.4:8466')
```
Launch TensorBoard and point it to the same logdir you provided to this API.
```shell
# logdir can be gs://your_tb_dir as in the above examples.
$ tensorboard --logdir=/tmp/tb_log
```
Open your browser and go to localhost:6006/#profile to view profiling results.
"""
if duration_ms <= 0:
raise errors.InvalidArgumentError(None, None,
'duration_ms must be greater than zero.')
opts = dict(options._asdict()) if options is not None else {}
_pywrap_profiler_plugin.trace(
_strip_addresses(service_addr, _GRPC_PREFIX),
logdir,
worker_list,
True,
duration_ms,
num_tracing_attempts,
opts,
)
@tf_export('profiler.experimental.client.monitor', v1=[])
def monitor(service_addr, duration_ms, level=1):
"""Sends grpc requests to profiler server to perform on-demand monitoring.
The monitoring result is a light weight performance summary of your model
execution. This method will block the caller thread until it receives the
monitoring result. This method currently supports Cloud TPU only.
Args:
service_addr: gRPC address of profiler service e.g. grpc://10.0.0.2:8466.
duration_ms: Duration of monitoring in ms.
level: Choose a monitoring level between 1 and 2 to monitor your job. Level
2 is more verbose than level 1 and shows more metrics.
Returns:
A string of monitoring output.
Example usage:
```python
# Continuously send gRPC requests to the Cloud TPU to monitor the model
# execution.
for query in range(0, 100):
print(
tf.profiler.experimental.client.monitor('grpc://10.0.0.2:8466', 1000))
```
"""
return _pywrap_profiler_plugin.monitor(
_strip_prefix(service_addr, _GRPC_PREFIX), duration_ms, level, True
)
def _strip_prefix(s, prefix):
return s[len(prefix):] if s.startswith(prefix) else s
def _strip_addresses(addresses, prefix):
return ','.join([_strip_prefix(s, prefix) for s in addresses.split(',')])
@@ -0,0 +1,68 @@
# Copyright 2020 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.
# ==============================================================================
"""Tests for profiler_client."""
import portpicker
from tensorflow.python.eager import test
from tensorflow.python.framework import test_util
from tensorflow.python.profiler import profiler_client
from tensorflow.python.profiler import profiler_v2 as profiler
class ProfilerClientTest(test_util.TensorFlowTestCase):
def testTrace_ProfileIdleServer(self):
test_port = portpicker.pick_unused_port()
profiler.start_server(test_port)
# Test the profilers are successfully started and connected to profiler
# service on the worker. Since there is no op running, it is expected to
# return RuntimeError with no trace events collected string.
with self.assertRaises(RuntimeError) as error:
profiler_client.trace(
f'localhost:{test_port}', self.get_temp_dir(), duration_ms=10
)
self.assertStartsWith(
str(error.exception), 'UNAVAILABLE: No trace event was collected'
)
def testTrace_ProfileIdleServerWithOptions(self):
test_port = portpicker.pick_unused_port()
profiler.start_server(test_port)
# Test the profilers are successfully started and connected to profiler
# service on the worker. Since there is no op running, it is expected to
# return RuntimeError with no trace events collected string.
with self.assertRaises(RuntimeError) as error:
options = profiler.ProfilerOptions(
host_tracer_level=3, device_tracer_level=0
)
profiler_client.trace(
f'localhost:{test_port}',
self.get_temp_dir(),
duration_ms=10,
options=options,
)
self.assertStartsWith(
str(error.exception), 'UNAVAILABLE: No trace event was collected'
)
def testMonitor_ProcessInvalidAddress(self):
# Monitor is only supported in cloud TPU. Test invalid address instead.
with self.assertRaises(RuntimeError):
profiler_client.monitor('localhost:6006', 2000)
if __name__ == '__main__':
test.main()
+220
View File
@@ -0,0 +1,220 @@
# Copyright 2016 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.
# ==============================================================================
import os
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.client import session
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import variables
from tensorflow.python.platform import gfile
from tensorflow.python.platform import test
from tensorflow.python.profiler import option_builder
# pylint: disable=g-bad-import-order
from tensorflow.python.profiler import model_analyzer
from tensorflow.python.profiler.internal import model_analyzer_testlib as lib
builder = option_builder.ProfileOptionBuilder
class ProfilerTest(test.TestCase):
@test_util.run_deprecated_v1
def testProfileBasic(self):
ops.reset_default_graph()
outfile = os.path.join(test.get_temp_dir(), 'dump')
opts = (builder(builder.trainable_variables_parameter())
.with_file_output(outfile)
.with_accounted_types(['.*'])
.select(['params', 'float_ops', 'micros', 'bytes',
'device', 'op_types', 'occurrence']).build())
# Test the output without run_meta.
sess = session.Session()
r = lib.BuildFullModel()
sess.run(variables.global_variables_initializer())
# Test the output with run_meta.
run_meta = config_pb2.RunMetadata()
_ = sess.run(r,
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE),
run_metadata=run_meta)
profiler = model_analyzer.Profiler(sess.graph)
profiler.add_step(1, run_meta)
profiler.profile_graph(opts)
with gfile.Open(outfile, 'r') as f:
profiler_str = f.read()
model_analyzer.profile(
sess.graph, cmd='graph', run_meta=run_meta, options=opts)
with gfile.Open(outfile, 'r') as f:
pma_str = f.read()
self.assertEqual(pma_str, profiler_str)
profiler.profile_name_scope(opts)
with gfile.Open(outfile, 'r') as f:
profiler_str = f.read()
model_analyzer.profile(
sess.graph, cmd='scope', run_meta=run_meta, options=opts)
with gfile.Open(outfile, 'r') as f:
pma_str = f.read()
self.assertEqual(pma_str, profiler_str)
profiler.profile_python(opts)
with gfile.Open(outfile, 'r') as f:
profiler_str = f.read()
model_analyzer.profile(
sess.graph, cmd='code', run_meta=run_meta, options=opts)
with gfile.Open(outfile, 'r') as f:
pma_str = f.read()
self.assertEqual(pma_str, profiler_str)
profiler.profile_operations(opts)
with gfile.Open(outfile, 'r') as f:
profiler_str = f.read()
model_analyzer.profile(
sess.graph, cmd='op', run_meta=run_meta, options=opts)
with gfile.Open(outfile, 'r') as f:
pma_str = f.read()
self.assertEqual(pma_str, profiler_str)
model_analyzer.profile(
sess.graph, cmd='scope', run_meta=run_meta, options=opts)
with gfile.Open(outfile, 'r') as f:
pma_str = f.read()
self.assertNotEqual(pma_str, profiler_str)
def testMultiStepProfile(self):
ops.reset_default_graph()
opts = builder.time_and_memory(min_bytes=0)
with session.Session() as sess:
r1, r2, r3 = lib.BuildSplittableModel()
sess.run(variables.global_variables_initializer())
profiler = model_analyzer.Profiler(sess.graph)
pb0 = profiler.profile_name_scope(opts)
run_meta = config_pb2.RunMetadata()
_ = sess.run(r1,
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE),
run_metadata=run_meta)
profiler.add_step(1, run_meta)
pb1 = profiler.profile_name_scope(opts)
self.assertNotEqual(lib.SearchTFProfNode(pb1, 'DW'), None)
self.assertEqual(lib.SearchTFProfNode(pb1, 'DW2'), None)
self.assertEqual(lib.SearchTFProfNode(pb1, 'add'), None)
run_meta2 = config_pb2.RunMetadata()
_ = sess.run(r2,
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE),
run_metadata=run_meta2)
profiler.add_step(2, run_meta2)
pb2 = profiler.profile_name_scope(opts)
self.assertNotEqual(lib.SearchTFProfNode(pb2, 'DW'), None)
self.assertNotEqual(lib.SearchTFProfNode(pb2, 'DW2'), None)
self.assertEqual(lib.SearchTFProfNode(pb2, 'add'), None)
run_meta3 = config_pb2.RunMetadata()
_ = sess.run(r3,
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE),
run_metadata=run_meta3)
profiler.add_step(3, run_meta3)
pb3 = profiler.profile_name_scope(opts)
self.assertNotEqual(lib.SearchTFProfNode(pb3, 'DW'), None)
self.assertNotEqual(lib.SearchTFProfNode(pb3, 'DW2'), None)
self.assertNotEqual(lib.SearchTFProfNode(pb3, 'add'), None)
self.assertEqual(lib.SearchTFProfNode(pb0, 'Conv2D'), None)
self.assertGreater(lib.SearchTFProfNode(pb1, 'Conv2D').exec_micros, 0)
self.assertEqual(lib.SearchTFProfNode(pb1, 'Conv2D_1'), None)
self.assertGreater(lib.SearchTFProfNode(pb2, 'Conv2D_1').exec_micros, 0)
self.assertEqual(lib.SearchTFProfNode(pb2, 'add'), None)
self.assertGreater(lib.SearchTFProfNode(pb3, 'add').exec_micros, 0)
advice_pb = profiler.advise(model_analyzer.ALL_ADVICE)
self.assertTrue('AcceleratorUtilizationChecker' in advice_pb.checkers)
self.assertTrue('ExpensiveOperationChecker' in advice_pb.checkers)
self.assertTrue('OperationChecker' in advice_pb.checkers)
checker = advice_pb.checkers['AcceleratorUtilizationChecker']
if test.is_gpu_available():
self.assertGreater(len(checker.reports), 0)
else:
self.assertEqual(len(checker.reports), 0)
checker = advice_pb.checkers['ExpensiveOperationChecker']
self.assertGreater(len(checker.reports), 0)
@test_util.run_deprecated_v1
def testMultipleProfilePerStep(self):
ops.reset_default_graph()
opts = (builder(builder.trainable_variables_parameter())
.with_empty_output()
.with_accounted_types(['.*'])
.select(['micros', 'bytes', 'peak_bytes',
'residual_bytes', 'output_bytes']).build())
r = lib.BuildSmallModel()
sess = session.Session()
profiler = model_analyzer.Profiler(sess.graph)
init_var_run_meta = config_pb2.RunMetadata()
sess.run(variables.global_variables_initializer(),
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE),
run_metadata=init_var_run_meta)
train_run_meta = config_pb2.RunMetadata()
sess.run(r,
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE),
run_metadata=train_run_meta)
profiler.add_step(0, train_run_meta)
ret1 = profiler.profile_name_scope(opts)
n1 = lib.SearchTFProfNode(
ret1, 'DW/Initializer/random_normal/RandomStandardNormal')
# Without the var initialization run_meta, it doesn't have the
# information of var_initialization.
self.assertEqual(n1.exec_micros, 0)
self.assertEqual(n1.requested_bytes, 0)
self.assertEqual(n1.peak_bytes, 0)
self.assertEqual(n1.residual_bytes, 0)
profiler.add_step(0, init_var_run_meta)
ret2 = profiler.profile_name_scope(opts)
n2 = lib.SearchTFProfNode(
ret2, 'DW/Initializer/random_normal/RandomStandardNormal')
# After adding the var initialization run_meta.
self.assertGreater(n2.exec_micros, 0)
self.assertGreater(n2.requested_bytes, 0)
self.assertGreater(n2.peak_bytes, 0)
self.assertGreater(n2.residual_bytes, 0)
if __name__ == '__main__':
test.main()
+212
View File
@@ -0,0 +1,212 @@
# Copyright 2020 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.
# ==============================================================================
"""TensorFlow 2.x Profiler.
The profiler has two modes:
- Programmatic Mode: start(logdir), stop(), and Profiler class. Profiling starts
when calling start(logdir) or create a Profiler class.
Profiling stops when calling stop() to save to
TensorBoard logdir or destroying the Profiler class.
- Sampling Mode: start_server(). It will perform profiling after receiving a
profiling request.
NOTE: Only one active profiler session is allowed. Use of simultaneous
Programmatic Mode and Sampling Mode is undefined and will likely fail.
NOTE: The Keras TensorBoard callback will automatically perform sampled
profiling. Before enabling customized profiling, set the callback flag
"profile_batches=[]" to disable automatic sampled profiling.
"""
import collections
import threading
from tensorflow.python.framework import errors
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.profiler.internal import _pywrap_profiler
from tensorflow.python.util.tf_export import tf_export
_profiler = None
_profiler_lock = threading.Lock()
@tf_export('profiler.experimental.ProfilerOptions', v1=[])
class ProfilerOptions(
collections.namedtuple('ProfilerOptions', [
'host_tracer_level', 'python_tracer_level', 'device_tracer_level',
'delay_ms'
])):
"""Options for finer control over the profiler.
Use `tf.profiler.experimental.ProfilerOptions` to control `tf.profiler`
behavior.
Fields:
host_tracer_level: Adjust CPU tracing level. Values are: `1` - critical info
only, `2` - info, `3` - verbose. [default value is `2`]
python_tracer_level: Toggle tracing of Python function calls. Values are:
`1` - enabled, `0` - disabled [default value is `0`]
device_tracer_level: Adjust device (TPU/GPU) tracing level. Values are:
`1` - enabled, `0` - disabled [default value is `1`]
delay_ms: Requests for all hosts to start profiling at a timestamp that is
`delay_ms` away from the current time. `delay_ms` is in milliseconds. If
zero, each host will start profiling immediately upon receiving the
request. Default value is `None`, allowing the profiler guess the best
value.
"""
def __new__(cls,
host_tracer_level=2,
python_tracer_level=0,
device_tracer_level=1,
delay_ms=None):
return super(ProfilerOptions,
cls).__new__(cls, host_tracer_level, python_tracer_level,
device_tracer_level, delay_ms)
@tf_export('profiler.experimental.start', v1=[])
def start(logdir, options=None):
"""Start profiling TensorFlow performance.
Args:
logdir: Profiling results log directory.
options: `ProfilerOptions` namedtuple to specify miscellaneous profiler
options. See example usage below.
Raises:
AlreadyExistsError: If a profiling session is already running.
Example usage:
```python
options = tf.profiler.experimental.ProfilerOptions(host_tracer_level = 3,
python_tracer_level = 1,
device_tracer_level = 1)
tf.profiler.experimental.start('logdir_path', options = options)
# Training code here
tf.profiler.experimental.stop()
```
To view the profiling results, launch TensorBoard and point it to `logdir`.
Open your browser and go to `localhost:6006/#profile` to view profiling
results.
"""
global _profiler
with _profiler_lock:
if _profiler is not None:
raise errors.AlreadyExistsError(None, None,
'Another profiler is running.')
_profiler = _pywrap_profiler.ProfilerSession()
try:
# support for namedtuple in pybind11 is missing, we change it to
# dict type first.
opts = dict(options._asdict()) if options is not None else {}
_profiler.start(logdir, opts)
except errors.AlreadyExistsError:
logging.warning('Another profiler session is running which is probably '
'created by profiler server. Please avoid using profiler '
'server and profiler APIs at the same time.')
raise errors.AlreadyExistsError(None, None,
'Another profiler is running.')
except Exception:
_profiler = None
raise
@tf_export('profiler.experimental.stop', v1=[])
def stop(save=True):
"""Stops the current profiling session.
The profiler session will be stopped and profile results can be saved.
Args:
save: An optional variable to save the results to TensorBoard. Default True.
Raises:
UnavailableError: If there is no active profiling session.
"""
global _profiler
with _profiler_lock:
if _profiler is None:
raise errors.UnavailableError(
None, None,
'Cannot export profiling results. No profiler is running.')
if save:
try:
_profiler.export_to_tb()
except Exception:
_profiler = None
raise
_profiler = None
def warmup():
"""Warm-up the profiler session.
The profiler session will set up profiling context, including loading CUPTI
library for GPU profiling. This is used for improving the accuracy of
the profiling results.
"""
start('')
stop(save=False)
@tf_export('profiler.experimental.server.start', v1=[])
def start_server(port):
"""Start a profiler grpc server that listens to given port.
The profiler server will exit when the process finishes. The service is
defined in tensorflow/core/profiler/profiler_service.proto.
Args:
port: port profiler server listens to.
Example usage: ```python tf.profiler.experimental.server.start(6009) # do
your training here.
"""
_pywrap_profiler.start_server(port)
@tf_export('profiler.experimental.Profile', v1=[])
class Profile(object):
"""Context-manager profile API.
Profiling will start when entering the scope, and stop and save the results to
the logdir when exits the scope. Open TensorBoard profile tab to view results.
Example usage:
```python
with tf.profiler.experimental.Profile("/path/to/logdir"):
# do some work
```
"""
def __init__(self, logdir, options=None):
"""Creates a context manager object for profiler API.
Args:
logdir: profile data will save to this directory.
options: An optional `tf.profiler.experimental.ProfilerOptions` can be
provided to fine tune the profiler's behavior.
"""
self._logdir = logdir
self._options = options
def __enter__(self):
start(self._logdir, self._options)
def __exit__(self, typ, value, tb):
stop()
@@ -0,0 +1,111 @@
# Copyright 2020 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.
# ==============================================================================
"""Tests for tf 2.x profiler."""
import os
from tensorflow.python.eager import test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.platform import gfile
from tensorflow.python.profiler import profiler_v2 as profiler
from tensorflow.python.profiler import trace
class ProfilerTest(test_util.TensorFlowTestCase):
def test_profile_exceptions(self):
logdir = self.get_temp_dir()
profiler.start(logdir)
with self.assertRaises(errors.AlreadyExistsError):
profiler.start(logdir)
profiler.stop()
with self.assertRaises(errors.UnavailableError):
profiler.stop()
# Test with a bad logdir, and it correctly raises exception and deletes
# profiler.
# pylint: disable=anomalous-backslash-in-string
profiler.start('/dev/null/\/\/:123')
# pylint: enable=anomalous-backslash-in-string
with self.assertRaises(Exception):
profiler.stop()
profiler.start(logdir)
profiler.stop()
def test_save_profile(self):
logdir = self.get_temp_dir()
profiler.start(logdir)
with trace.Trace('three_times_five'):
three = constant_op.constant(3)
five = constant_op.constant(5)
product = three * five
self.assertAllEqual(15, product)
profiler.stop()
file_list = gfile.ListDirectory(logdir)
self.assertEqual(len(file_list), 1)
for file_name in gfile.ListDirectory(logdir):
if gfile.IsDirectory(os.path.join(logdir, file_name)):
self.assertEqual(file_name, 'plugins')
profile_dir = os.path.join(logdir, 'plugins', 'profile')
run = gfile.ListDirectory(profile_dir)[0]
run_dir = os.path.join(profile_dir, run)
xplanes = gfile.Glob(os.path.join(run_dir, '*.xplane.pb'))
self.assertLen(xplanes, 1)
self.assertTrue(gfile.Exists(xplanes[0]))
def test_profile_with_options(self):
logdir = self.get_temp_dir()
options = profiler.ProfilerOptions(
host_tracer_level=3, python_tracer_level=1)
profiler.start(logdir, options)
with trace.Trace('three_times_five'):
three = constant_op.constant(3)
five = constant_op.constant(5)
product = three * five
self.assertAllEqual(15, product)
profiler.stop()
file_list = gfile.ListDirectory(logdir)
self.assertEqual(len(file_list), 1)
def test_context_manager_with_options(self):
logdir = self.get_temp_dir()
options = profiler.ProfilerOptions(
host_tracer_level=3, python_tracer_level=1)
with profiler.Profile(logdir, options):
with trace.Trace('three_times_five'):
three = constant_op.constant(3)
five = constant_op.constant(5)
product = three * five
self.assertAllEqual(15, product)
file_list = gfile.ListDirectory(logdir)
self.assertEqual(len(file_list), 1)
def test_callback(self):
logdir = self.get_temp_dir()
self.assertFalse(trace.enabled())
profiler.start(logdir)
self.assertTrue(trace.enabled())
profiler.stop()
self.assertFalse(trace.enabled())
if __name__ == '__main__':
test.main()
@@ -0,0 +1,41 @@
# Copyright 2020 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.
# ==============================================================================
"""Tests for profiler_wrapper.cc pybind methods."""
from tensorflow.python.eager import test
from tensorflow.python.framework import test_util
from tensorflow.python.profiler.internal import _pywrap_profiler_plugin as profiler_wrapper_plugin
class ProfilerSessionTest(test_util.TensorFlowTestCase):
def test_xspace_to_tools_data_default_options(self):
# filenames only used for `hlo_proto` tool.
profiler_wrapper_plugin.xspace_to_tools_data([], 'trace_viewer')
def _test_xspace_to_tools_data_options(self, options):
profiler_wrapper_plugin.xspace_to_tools_data([], 'trace_viewer', options)
def test_xspace_to_tools_data_empty_options(self):
self._test_xspace_to_tools_data_options({})
def test_xspace_to_tools_data_int_options(self):
self._test_xspace_to_tools_data_options({'example_option': 0})
def test_xspace_to_tools_data_str_options(self):
self._test_xspace_to_tools_data_options({'example_option': 'example'})
if __name__ == '__main__':
test.main()
+215
View File
@@ -0,0 +1,215 @@
# 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.
# ==============================================================================
"""Logging tensorflow::tfprof::OpLogProto.
OpLogProto is used to add extra model information for offline analysis.
"""
import os
import sys
from tensorflow.core.profiler import tfprof_log_pb2
from tensorflow.python.eager import context
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.platform import gfile
from tensorflow.python.profiler.internal import flops_registry # pylint: disable=unused-import
from tensorflow.python.util.tf_export import tf_export
TRAINABLE_VARIABLES = '_trainable_variables'
REGISTERED_FLOP_STATS = 'flops'
def _fill_missing_graph_shape(graph, run_meta):
"""Fill Tensor shapes in 'graph' with run time shape from 'run_meta'."""
for dev_stat in run_meta.step_stats.dev_stats:
for node_stat in dev_stat.node_stats:
if not node_stat.output:
continue
try:
op = graph.get_operation_by_name(node_stat.node_name)
except KeyError as e:
# Graph doesn't contains the node_stat, usually RecvTensor.
continue
if len(node_stat.output) != len(op.outputs):
# For example, conditional op has only 1 output at run time.
continue
for (i, node_stat_out) in enumerate(node_stat.output):
if op.outputs[i].get_shape().is_fully_defined():
continue
node_stat_dims = node_stat_out.tensor_description.shape.dim
node_stat_shape = tensor_shape.TensorShape(
[d.size for d in node_stat_dims])
try:
op.outputs[i].set_shape(op.outputs[i].get_shape().merge_with(
node_stat_shape))
except ValueError as e:
sys.stderr.write('Node %s incompatible shapes: %s.\n' %
(node_stat.node_name, e))
return graph
def _str_id(s, str_to_id):
"""Maps string to id."""
num = str_to_id.get(s, None)
if num is None:
num = len(str_to_id)
str_to_id[s] = num
return num
def _get_logged_ops(graph, run_meta=None, add_trace=True,
add_trainable_var=True):
"""Extract trainable model parameters and FLOPs for ops from a Graph.
Args:
graph: tf.Graph.
run_meta: RunMetadata proto used to complete shape information.
add_trace: Whether to add op trace information.
add_trainable_var: Whether to assign tf.compat.v1.trainable_variables() op
type '_trainable_variables'.
Returns:
logged_ops: dict mapping from op_name to OpLogEntry.
string_to_id: dict mapping from string to id.
"""
if run_meta:
graph = _fill_missing_graph_shape(graph, run_meta)
op_missing_shape = 0
logged_ops = {}
string_to_id = {}
string_to_id['none'] = len(string_to_id)
# TODO(xpan): Work with Profiler more efficiently.
for op in graph.get_operations():
try:
stats = ops.get_stats_for_node_def(
graph, op.node_def, REGISTERED_FLOP_STATS)
except ValueError:
# Catch Exception When shape is incomplete. Skip it.
op_missing_shape += 1
stats = None
entry = tfprof_log_pb2.OpLogEntry()
entry.name = op.name
add_entry = False
if stats and stats.value:
entry.float_ops = int(stats.value)
add_entry = True
if add_trace:
if op.traceback:
for filename, lineno, funcname, line in op.traceback:
trace = entry.code_def.traces.add()
trace.file_id = _str_id(filename, string_to_id) if filename else 0
trace.lineno = lineno if lineno else -1
trace.function_id = _str_id(funcname, string_to_id) if funcname else 0
trace.line_id = _str_id(line, string_to_id) if line else 0
# TODO(slebedev): remove this unused field from the proto.
trace.func_start_line = -1
add_entry = True
if add_entry:
logged_ops[entry.name] = entry
if add_trainable_var:
for v in graph.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES):
if v.op.name not in logged_ops:
entry = tfprof_log_pb2.OpLogEntry()
entry.name = v.op.name
entry.types.append(TRAINABLE_VARIABLES)
logged_ops[entry.name] = entry
else:
logged_ops[v.op.name].types.append(TRAINABLE_VARIABLES)
if op_missing_shape > 0 and not run_meta:
sys.stderr.write('%d ops no flops stats due to incomplete shapes.\n' %
op_missing_shape)
return logged_ops, string_to_id
def merge_default_with_oplog(graph, op_log=None, run_meta=None,
add_trace=True, add_trainable_var=True):
"""Merge the tfprof default extra info with caller's op_log.
Args:
graph: tf.Graph. If None and eager execution is not enabled, use
default graph.
op_log: OpLogProto proto.
run_meta: RunMetadata proto used to complete shape information.
add_trace: Whether to add op trace information.
add_trainable_var: Whether to assign tf.compat.v1.trainable_variables() op
type '_trainable_variables'.
Returns:
tmp_op_log: Merged OpLogProto proto.
"""
if not graph and not context.executing_eagerly():
graph = ops.get_default_graph()
tmp_op_log = tfprof_log_pb2.OpLogProto()
if not graph:
return tmp_op_log
logged_ops, string_to_id = _get_logged_ops(
graph, run_meta, add_trace=add_trace, add_trainable_var=add_trainable_var)
if not op_log:
tmp_op_log.log_entries.extend(logged_ops.values())
else:
all_ops = {}
for entry in op_log.log_entries:
all_ops[entry.name] = entry
for op_name, entry in logged_ops.items():
if op_name in all_ops:
all_ops[op_name].types.extend(entry.types)
if entry.float_ops > 0 and all_ops[op_name].float_ops == 0:
all_ops[op_name].float_ops = entry.float_ops
if entry.code_def.traces and not all_ops[op_name].code_def.traces:
all_ops[op_name].code_def.MergeFrom(entry.code_def)
else:
all_ops[op_name] = entry
tmp_op_log.log_entries.extend(all_ops.values())
for s, i in string_to_id.items():
tmp_op_log.id_to_string[i] = s
return tmp_op_log
@tf_export(v1=['profiler.write_op_log'])
def write_op_log(graph, log_dir, op_log=None, run_meta=None, add_trace=True):
"""Log provided 'op_log', and add additional model information below.
The API also assigns ops in tf.compat.v1.trainable_variables() an op type
called '_trainable_variables'.
The API also logs 'flops' statistics for ops with op.RegisterStatistics()
defined. flops calculation depends on Tensor shapes defined in 'graph',
which might not be complete. 'run_meta', if provided, completes the shape
information with best effort.
Args:
graph: tf.Graph. If None and eager execution is not enabled, use
default graph.
log_dir: directory to write the log file.
op_log: (Optional) OpLogProto proto to be written. If not provided, an new
one is created.
run_meta: (Optional) RunMetadata proto that helps flops computation using
run time shape information.
add_trace: Whether to add python code trace information.
Used to support "code" view.
"""
if not graph and not context.executing_eagerly():
graph = ops.get_default_graph()
op_log = merge_default_with_oplog(graph, op_log, run_meta, add_trace)
with gfile.Open(os.path.join(log_dir, 'tfprof_log'), 'w') as log:
log.write(op_log.SerializeToString())
@@ -0,0 +1,76 @@
# Copyright 2016 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.
# ==============================================================================
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
class TFProfLoggerTest(test.TestCase):
def _BuildSmallPlaceholderlModel(self):
a = array_ops.placeholder(dtypes.int32, [2, 2])
b = array_ops.placeholder(dtypes.int32, [2, 2])
y = math_ops.matmul(a, b)
return a, b, y
def _BuildSmallModel(self):
a = constant_op.constant([[1, 2], [3, 4]])
b = constant_op.constant([[1, 2], [3, 4]])
return math_ops.matmul(a, b)
# pylint: disable=pointless-string-statement
"""# TODO(xpan): This out of core so it doesn't depend on contrib.
def testFillMissingShape(self):
a, b, y = self._BuildSmallPlaceholderlModel()
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
run_metadata = config_pb2.RunMetadata()
sess = session.Session()
sess.run(y,
options=run_options,
run_metadata=run_metadata,
feed_dict={a: [[1, 2], [2, 3]],
b: [[1, 2], [2, 3]]})
graph2 = ops.Graph()
# Use copy_op_to_graph to remove shape information.
y2 = copy_elements.copy_op_to_graph(y, graph2, [])
self.assertEqual('<unknown>', str(y2.get_shape()))
tfprof_logger._fill_missing_graph_shape(graph2, run_metadata)
self.assertEqual('(2, 2)', str(y2.get_shape()))
def testFailedFillMissingShape(self):
y = self._BuildSmallModel()
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
run_metadata = config_pb2.RunMetadata()
sess = session.Session()
sess.run(y, options=run_options, run_metadata=run_metadata)
graph2 = ops.Graph()
y2 = copy_elements.copy_op_to_graph(y, graph2, [])
self.assertEqual('<unknown>', str(y2.get_shape()))
# run_metadata has special name for MatMul, hence failed to fill shape.
tfprof_logger._fill_missing_graph_shape(graph2, run_metadata)
self.assertEqual('<unknown>', str(y2.get_shape()))
"""
if __name__ == '__main__':
test.main()
+187
View File
@@ -0,0 +1,187 @@
# 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.
# ==============================================================================
"""Trace allows the profiler to trace Python events."""
import functools
from tensorflow.python.profiler.internal import _pywrap_traceme
from tensorflow.python.util.tf_export import tf_export
# This is a low-overhead function that directly calls C++ to check if the
# profiler is enabled.
enabled = _pywrap_traceme.traceme_enabled
@tf_export('profiler.experimental.Trace', v1=[])
class Trace(object):
"""Context manager that generates a trace event in the profiler.
A trace event will start when entering the context, and stop and save the
result to the profiler when exiting the context. Open TensorBoard Profile tab
and choose trace viewer to view the trace event in the timeline.
Trace events are created only when the profiler is enabled. More information
on how to use the profiler can be found at
https://tensorflow.org/guide/profiler
Example usage:
```python
tf.profiler.experimental.start('logdir')
for step in range(num_steps):
# Creates a trace event for each training step with the step number.
with tf.profiler.experimental.Trace("Train", step_num=step, _r=1):
train_fn()
tf.profiler.experimental.stop()
```
"""
def __init__(self, name, **kwargs):
"""Creates a trace event in the profiler.
Args:
name: The name of the trace event.
**kwargs: Keyword arguments added to the trace event.
Both the key and value are of types that
can be converted to strings, which will be
interpreted by the profiler according to the
traceme name.
Example usage:
```python
tf.profiler.experimental.start('logdir')
for step in range(num_steps):
# Creates a trace event for each training step with the
# step number.
with tf.profiler.experimental.Trace("Train", step_num=step):
train_fn()
tf.profiler.experimental.stop()
```
The example above uses the keyword argument "step_num" to specify the
training step being traced.
"""
if enabled():
# Creating _pywrap_traceme.TraceMe starts the clock.
self._traceme = _pywrap_traceme.TraceMe(name, **kwargs)
else:
self._traceme = None
def __enter__(self):
# Starting the TraceMe clock here would require an extra Python->C++ call.
return self
def set_metadata(self, **kwargs):
"""Sets metadata in this trace event.
Args:
**kwargs: metadata in key-value pairs.
This method enables setting metadata in a trace event after it is
created.
Example usage:
```python
def call(function):
with tf.profiler.experimental.Trace("call",
function_name=function.name) as tm:
binary, in_cache = jit_compile(function)
tm.set_metadata(in_cache=in_cache)
execute(binary)
```
In this example, we want to trace how much time spent on
calling a function, which includes compilation and execution.
The compilation can be either getting a cached copy of the
binary or actually generating the binary, which is indicated
by the boolean "in_cache" returned by jit_compile(). We need
to use set_metadata() to pass in_cache because we did not know
the in_cache value when the trace was created (and we cannot
create the trace after jit_compile(), because we want
to measure the entire duration of call()).
"""
if self._traceme and kwargs:
self._traceme.SetMetadata(**kwargs)
def __exit__(self, exc_type, exc_val, exc_tb):
if self._traceme:
self._traceme.Stop()
def trace_wrapper(trace_name, **trace_kwargs):
"""Decorator alternative to `with Trace(): ...`. It's faster.
Args:
trace_name: The name of the trace event, or a callable to be traced, in
which case the name is inferred from qualname or name of the callable.
**trace_kwargs: Keyword arguments added to the trace event. Both the key and
value are of types that can be converted to strings, which will be
interpreted by the profiler according to the traceme name.
Returns:
A decorator that can wrap a function and apply `Trace` scope if needed,
or a decorated function if used as a decorator directly.
Example usage:
```python
@trace_wrapper('trace_name')
def func(x, y, z):
pass # code to execute and apply `Trace` if needed.
# Equivalent to
# with Trace('trace_name'):
# func(1, 2, 3)
func(1, 2, 3)
```
or
```python
@trace_wrapper
def func(x, y, z):
pass # code to execute and apply `Trace` if needed.
# Equivalent to
# with Trace(func.__qualname__):
# func(1, 2, 3)
func(1, 2, 3)
```
"""
if callable(trace_name):
func = trace_name
name = getattr(func, '__qualname__', None)
if not name:
name = getattr(func, '__name__', 'unknown function')
return trace_wrapper(name)(func)
def inner_wrapper(func):
@functools.wraps(func)
def wrapped(*args, **kwargs):
if enabled():
with Trace(trace_name, **trace_kwargs):
return func(*args, **kwargs)
return func(*args, **kwargs)
return wrapped
return inner_wrapper