Revision f70e2dc8ff39f097fcd102407941fc1fb5e347ea authored by Alexander G. de G. Matthews on 04 December 2017, 19:33:04 UTC, committed by GitHub on 04 December 2017, 19:33:04 UTC
session_manager.py
# Copyright 2017 Artem Artemev @awav
#
# 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
import warnings
import tensorflow as tf
from tensorflow.python.client import timeline
from . import settings
class _DefaultSessionKeeper:
session = None
class TracerSession(tf.Session):
def __init__(self, output_file_name=None, output_directory=None,
each_time=None, **kwargs):
self.output_file_name = output_file_name
self.output_directory = output_directory
self.each_time = each_time
self.local_run_metadata = None
if self.each_time:
warnings.warn("Outputting a trace for each run. "
"May result in large disk usage.")
super(TracerSession, self).__init__(**kwargs)
self.counter = 0
self.profiler_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
if self.output_directory is not None:
if os.path.isfile(self.output_directory):
raise IOError("In tracer: given directory name is a file.")
if not os.path.isdir(self.output_directory):
os.mkdir(self.output_directory)
def _trace_filename(self):
"""
Creates trace filename.
"""
dir_stub = ''
if self.output_directory is not None:
dir_stub = self.output_directory
if self.each_time:
filename = '{0}_{1}.json'.format(
self.output_file_name, self.counter)
else:
filename = '{0}.json'.format(self.output_file_name)
return os.path.join(dir_stub, filename)
def run(self, fetches, feed_dict=None, options=None, run_metadata=None):
# Make sure there is no disagreement doing this.
if options is not None:
if options.trace_level != self.profiler_options.trace_level: # pragma: no cover
raise ValueError(
'In profiler session. Inconsistent trace '
'level from run call') # pragma: no cover
self.profiler_options.update(options) # pragma: no cover
self.local_run_metadata = tf.RunMetadata()
output = super(TracerSession, self).run(
fetches, feed_dict=feed_dict,
options=self.profiler_options,
run_metadata=self.local_run_metadata)
trace_time = timeline.Timeline(self.local_run_metadata.step_stats)
ctf = trace_time.generate_chrome_trace_format()
with open(self._trace_filename(), 'w') as trace_file:
trace_file.write(ctf)
if self.each_time:
self.counter += 1
return output
def reset_default_session(*args, **kwargs):
_DefaultSessionKeeper.session = get_session(*args, **kwargs)
def get_default_session(*args, **kwargs):
reset = kwargs.pop('reset', False)
if reset or _DefaultSessionKeeper.session is None:
_DefaultSessionKeeper.session = get_session(*args, **kwargs)
return _DefaultSessionKeeper.session
def get_session(*args, **kwargs):
"""
Pass session configuration options
"""
if 'config' not in kwargs:
kwargs['config'] = tf.ConfigProto(**settings.session)
if settings.profiling.dump_timeline:
def fill_kwargs(key, value):
"""
Internal function for filling default None values with meaningful
values from gpflow settings.
"""
if kwargs.get(key) is None:
kwargs[key] = value
fill_kwargs('output_file_name', settings.profiling.output_file_name)
fill_kwargs('output_directory', settings.profiling.output_directory)
fill_kwargs('each_time', settings.profiling.each_time)
return TracerSession(*args, **kwargs)
kwargs.pop("output_file_name", None)
kwargs.pop("output_directory", None)
kwargs.pop("each_time", None)
return tf.Session(*args, **kwargs)
![swh spinner](/static/img/swh-spinner.gif)
Computing file changes ...