experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add a method for getting in-memory profiles from ProfileContexts

Browse Source 
PiperOrigin-RevId: 205434648
This commit is contained in:
A. Unique TensorFlower 2018-07-20 12:15:16 -07:00 committed by TensorFlower Gardener
parent 41b93403ac
commit 1d3440dea0
2 changed files with 31 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
tensorflow/python/profiler/profile_context.py
View File

@ -88,16 +88,19 @@ def _profiled_run(self,
to_profiles = self.profile_context._profile_candidates() to_profiles = self.profile_context._profile_candidates()
for to_prof in to_profiles: for to_prof in to_profiles:
cmd, opts, _ = to_prof cmd, opts, _ = to_prof
saved_views = self.profile_context._views.setdefault(cmd, {})
if self.profile_context._debug: if self.profile_context._debug:
sys.stderr.write('debug: profiling %s step: %d\n' % (cmd, step)) sys.stderr.write('debug: profiling %s step: %d\n' % (cmd, step))
if cmd == 'graph': if cmd == 'graph':
self.profile_context.profiler.profile_graph(opts) saved_views[step] = self.profile_context.profiler.profile_graph(opts)
elif cmd == 'scope': elif cmd == 'scope':
self.profile_context.profiler.profile_name_scope(opts) saved_views[step] = self.profile_context.profiler.profile_name_scope(
opts)
elif cmd == 'op': elif cmd == 'op':
self.profile_context.profiler.profile_operations(opts) saved_views[step] = self.profile_context.profiler.profile_operations(
opts)
elif cmd == 'code': elif cmd == 'code':
self.profile_context.profiler.profile_python(opts) saved_views[step] = self.profile_context.profiler.profile_python(opts)
else: else:
raise ValueError('Unknown cmd: %s\n' % cmd) raise ValueError('Unknown cmd: %s\n' % cmd)
return ret return ret
@ -185,8 +188,30 @@ class ProfileContext(object):
self._traced_steps = 0 self._traced_steps = 0
self._auto_profiles = [] self._auto_profiles = []
self._profiler = None self._profiler = None
self._views = {}
self._lock = threading.Lock() 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): def add_auto_profiling(self, cmd, options, profile_steps):
"""Traces and profiles at some session run steps. """Traces and profiles at some session run steps.

2
tensorflow/python/profiler/profile_context_test.py
View File

@ -61,6 +61,8 @@ class ProfilerContextTest(test.TestCase):
profile_str = f.read() profile_str = f.read()
gfile.Remove(outfile) gfile.Remove(outfile)
self.assertEqual(set([15, 50, 100]), set(pctx.get_profiles("op").keys()))
with lib.ProfilerFromFile( with lib.ProfilerFromFile(
os.path.join(test.get_temp_dir(), "profile_100")) as profiler: os.path.join(test.get_temp_dir(), "profile_100")) as profiler:
profiler.profile_operations(options=opts) profiler.profile_operations(options=opts)