diff --git a/tensorboard/backend/event_processing/event_file_loader.py b/tensorboard/backend/event_processing/event_file_loader.py index 28977f6d31..d1374b2414 100644 --- a/tensorboard/backend/event_processing/event_file_loader.py +++ b/tensorboard/backend/event_processing/event_file_loader.py @@ -21,6 +21,7 @@ import inspect from tensorboard.compat import tf +from tensorboard.compat import _pywrap_tensorflow from tensorboard.compat.proto import event_pb2 from tensorboard.util import platform_util from tensorboard.util import tb_logging @@ -38,7 +39,7 @@ def __init__(self, file_path): file_path = platform_util.readahead_file_path(file_path) logger.debug('Opening a record reader pointing at %s', file_path) with tf.errors.raise_exception_on_not_ok_status() as status: - self._reader = tf.compat.v1.pywrap_tensorflow.PyRecordReader_New( + self._reader = _pywrap_tensorflow.PyRecordReader_New( tf.compat.as_bytes(file_path), 0, tf.compat.as_bytes(''), status) # Store it for logging purposes. self._file_path = file_path diff --git a/tensorboard/compat/__init__.py b/tensorboard/compat/__init__.py index 06f92b925c..634f252d4c 100644 --- a/tensorboard/compat/__init__.py +++ b/tensorboard/compat/__init__.py @@ -72,3 +72,34 @@ def tf2(): # As a fallback, try `tensorflow.compat.v2` if it's defined. return tf.compat.v2 raise ImportError('cannot import tensorflow 2.0 API') + + +# TODO(https://github.com/tensorflow/tensorboard/issues/1711): remove this +@_lazy.lazy_load('tensorboard.compat._pywrap_tensorflow') +def _pywrap_tensorflow(): + """Provide pywrap_tensorflow access in TensorBoard. + + pywrap_tensorflow cannot be accessed from tf.python.pywrap_tensorflow + and needs to be imported using + `from tensorflow.python import pywrap_tensorflow`. Therefore, we provide + a separate accessor function for it here. + + NOTE: pywrap_tensorflow is not part of TensorFlow API and this + dependency will go away soon. + + Returns: + pywrap_tensorflow import, if available. + + Raises: + ImportError: if we couldn't import pywrap_tensorflow. + """ + try: + from tensorboard.compat import notf # pylint: disable=g-import-not-at-top + except ImportError: + try: + from tensorflow.python import pywrap_tensorflow # pylint: disable=g-import-not-at-top + return pywrap_tensorflow + except ImportError: + pass + from tensorboard.compat.tensorflow_stub import pywrap_tensorflow # pylint: disable=g-import-not-at-top + return pywrap_tensorflow diff --git a/tensorboard/compat/tensorflow_stub/__init__.py b/tensorboard/compat/tensorflow_stub/__init__.py index d08dd23f4b..e45d3c99f7 100644 --- a/tensorboard/compat/tensorflow_stub/__init__.py +++ b/tensorboard/compat/tensorflow_stub/__init__.py @@ -35,8 +35,5 @@ from . import pywrap_tensorflow # noqa from . import tensor_shape # noqa -# Set pywrap_tensorflow on v1 and avoid cycles on some imports -compat.v1.pywrap_tensorflow = pywrap_tensorflow - # Set a fake __version__ to help distinguish this as our own stub API. __version__ = 'stub' diff --git a/tensorboard/plugins/debugger/debugger_plugin_testlib.py b/tensorboard/plugins/debugger/debugger_plugin_testlib.py index 74a3242883..2f50bb533f 100644 --- a/tensorboard/plugins/debugger/debugger_plugin_testlib.py +++ b/tensorboard/plugins/debugger/debugger_plugin_testlib.py @@ -25,6 +25,7 @@ import portpicker # pylint: disable=import-error import tensorflow as tf +from tensorflow.python import pywrap_tensorflow from werkzeug import wrappers from werkzeug import test as werkzeug_test @@ -60,7 +61,7 @@ def setUp(self): self.log_dir = self.get_temp_dir() file_prefix = tf.compat.as_bytes( os.path.join(self.log_dir, 'events.debugger')) - writer = tf.compat.v1.pywrap_tensorflow.EventsWriter(file_prefix) + writer = pywrap_tensorflow.EventsWriter(file_prefix) device_name = '/job:localhost/replica:0/task:0/cpu:0' writer.WriteEvent( self._CreateEventWithDebugNumericSummary( @@ -106,7 +107,7 @@ def setUp(self): os.mkdir(run_foo_directory) file_prefix = tf.compat.as_bytes( os.path.join(run_foo_directory, 'events.debugger')) - writer = tf.compat.v1.pywrap_tensorflow.EventsWriter(file_prefix) + writer = pywrap_tensorflow.EventsWriter(file_prefix) writer.WriteEvent( self._CreateEventWithDebugNumericSummary( device_name=device_name, diff --git a/tensorboard/plugins/debugger/events_writer_manager.py b/tensorboard/plugins/debugger/events_writer_manager.py index 129c085420..86e6373aa1 100644 --- a/tensorboard/plugins/debugger/events_writer_manager.py +++ b/tensorboard/plugins/debugger/events_writer_manager.py @@ -27,6 +27,7 @@ import time import tensorflow as tf +from tensorflow.python import pywrap_tensorflow from tensorboard.util import tb_logging logger = tb_logging.get_logger() @@ -208,7 +209,7 @@ def _create_events_writer(self, directory): os.path.join(directory, DEBUGGER_EVENTS_FILE_STARTING_TEXT), time.time(), self._events_file_count) logger.info("Creating events file %s", file_path) - return tf.compat.v1.pywrap_tensorflow.EventsWriter(tf.compat.as_bytes(file_path)) + return pywrap_tensorflow.EventsWriter(tf.compat.as_bytes(file_path)) def _fetch_events_files_on_disk(self): """Obtains the names of debugger-related events files within the directory. diff --git a/tensorboard/plugins/projector/projector_plugin.py b/tensorboard/plugins/projector/projector_plugin.py index 03c641d5a8..1cea3af83f 100644 --- a/tensorboard/plugins/projector/projector_plugin.py +++ b/tensorboard/plugins/projector/projector_plugin.py @@ -32,6 +32,7 @@ from tensorboard.backend.http_util import Respond from tensorboard.compat import tf +from tensorboard.compat import _pywrap_tensorflow from tensorboard.plugins import base_plugin from tensorboard.plugins.projector.projector_config_pb2 import ProjectorConfig from tensorboard.util import tb_logging @@ -423,8 +424,7 @@ def _get_reader_for_run(self, run): reader = None if config.model_checkpoint_path and _using_tf(): try: - reader = tf.compat.v1.pywrap_tensorflow.NewCheckpointReader( - config.model_checkpoint_path) + reader = tf.train.load_checkpoint(config.model_checkpoint_path) except Exception: # pylint: disable=broad-except logger.warn('Failed reading "%s"', config.model_checkpoint_path) self.readers[run] = reader