diff options
author | A. Unique TensorFlower <gardener@tensorflow.org> | 2018-07-20 12:15:16 -0700 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-07-20 12:18:17 -0700 |
commit | 1d3440dea001d7466f6e9f40ddd3afdf94ed8dc4 (patch) | |
tree | 103f54e6c1f218de7a6b741a8fd6cec0aaec69d3 /tensorflow/python/profiler | |
parent | 41b93403ac8148bd880c749165c40840ddb95b44 (diff) |
Add a method for getting in-memory profiles from ProfileContexts
PiperOrigin-RevId: 205434648
Diffstat (limited to 'tensorflow/python/profiler')
-rw-r--r-- | tensorflow/python/profiler/profile_context.py | 33 | ||||
-rw-r--r-- | tensorflow/python/profiler/profile_context_test.py | 2 |
2 files changed, 31 insertions, 4 deletions
diff --git a/tensorflow/python/profiler/profile_context.py b/tensorflow/python/profiler/profile_context.py index 18eb66ef98..fa4260a712 100644 --- a/tensorflow/python/profiler/profile_context.py +++ b/tensorflow/python/profiler/profile_context.py @@ -88,16 +88,19 @@ def _profiled_run(self, 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': - self.profile_context.profiler.profile_graph(opts) + saved_views[step] = self.profile_context.profiler.profile_graph(opts) 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': - self.profile_context.profiler.profile_operations(opts) + saved_views[step] = self.profile_context.profiler.profile_operations( + opts) elif cmd == 'code': - self.profile_context.profiler.profile_python(opts) + saved_views[step] = self.profile_context.profiler.profile_python(opts) else: raise ValueError('Unknown cmd: %s\n' % cmd) return ret @@ -185,8 +188,30 @@ class ProfileContext(object): 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. diff --git a/tensorflow/python/profiler/profile_context_test.py b/tensorflow/python/profiler/profile_context_test.py index a623beee23..107ad443c3 100644 --- a/tensorflow/python/profiler/profile_context_test.py +++ b/tensorflow/python/profiler/profile_context_test.py @@ -61,6 +61,8 @@ class ProfilerContextTest(test.TestCase): 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) |