Pull legacy contrib behavior out of ResourceSummaryWriter
This refactors ResourceSummaryWriter to remove the `v2` parameter, which previously controlled whether the behavior reflected `create_summary_writer()`, aka the legacy tf.contrib.summary API, or `create_summary_writer_v2()`, the API exposed in TF 2. Instead, we just split the behavior into two distinct implementations, `_LegacyResourceSummaryWriter` for contrib and `_ResourceSummaryWriter` for TF 2. This makes each implementation simpler to read in advance of further changes to this class hierarchy.
The symbols have been given leading underscores to discourage their direct usage (since users should refer to the `SummaryWriter` interface only), and I've updated `NoopSummaryWriter` to `_NoopSummaryWriter` accordingly.
PiperOrigin-RevId: 366406490
Change-Id: I67a5b6b9a7ad90d13cd59db0991fe7b5c876e425
diff --git a/tensorflow/python/ops/summary_ops_v2.py b/tensorflow/python/ops/summary_ops_v2.py
index ec483e9..d6c23e8 100644
--- a/tensorflow/python/ops/summary_ops_v2.py
+++ b/tensorflow/python/ops/summary_ops_v2.py
@@ -299,16 +299,14 @@
raise NotImplementedError()
-class ResourceSummaryWriter(SummaryWriter):
+class _ResourceSummaryWriter(SummaryWriter):
"""Implementation of SummaryWriter using a SummaryWriterInterface resource."""
- def __init__(self, shared_name, init_op_fn, name=None, v2=False):
+ def __init__(self, shared_name, init_op_fn, name=None):
self._resource = gen_summary_ops.summary_writer(
shared_name=shared_name, name=name)
- # TODO(nickfelt): cache other constructed ops in graph mode
self._init_op_fn = init_op_fn
self._init_op = init_op_fn(self._resource)
- self._v2 = v2
self._closed = False
if context.executing_eagerly():
self._resource_deleter = resource_variable_ops.EagerResourceDeleter(
@@ -318,46 +316,72 @@
def set_as_default(self, step=None):
"""See `SummaryWriter.set_as_default`."""
- if self._v2 and context.executing_eagerly() and self._closed:
+ if context.executing_eagerly() and self._closed:
raise RuntimeError("SummaryWriter is already closed")
super().set_as_default(step)
def as_default(self, step=None):
"""See `SummaryWriter.as_default`."""
- if self._v2 and context.executing_eagerly() and self._closed:
+ if context.executing_eagerly() and self._closed:
raise RuntimeError("SummaryWriter is already closed")
return super().as_default(step)
def init(self):
"""See `SummaryWriter.init`."""
- if self._v2:
- if context.executing_eagerly() and self._closed:
- raise RuntimeError("SummaryWriter is already closed")
- return self._init_op
- # Legacy behavior allows re-initializing the resource.
- return self._init_op_fn(self._resource)
+ if context.executing_eagerly() and self._closed:
+ raise RuntimeError("SummaryWriter is already closed")
+ return self._init_op
def flush(self):
"""See `SummaryWriter.flush`."""
- if self._v2 and context.executing_eagerly() and self._closed:
+ if context.executing_eagerly() and self._closed:
return
with ops.device("cpu:0"):
return gen_summary_ops.flush_summary_writer(self._resource)
def close(self):
"""See `SummaryWriter.close`."""
- if self._v2 and context.executing_eagerly() and self._closed:
+ if context.executing_eagerly() and self._closed:
return
try:
with ops.control_dependencies([self.flush()]):
with ops.device("cpu:0"):
return gen_summary_ops.close_summary_writer(self._resource)
finally:
- if self._v2 and context.executing_eagerly():
+ if context.executing_eagerly():
self._closed = True
-class NoopSummaryWriter(SummaryWriter):
+class _LegacyResourceSummaryWriter(SummaryWriter):
+ """Legacy resource-backed SummaryWriter for tf.contrib.summary."""
+
+ def __init__(self, resource, init_op_fn):
+ self._resource = resource
+ self._init_op_fn = init_op_fn
+ init_op = self.init()
+ if context.executing_eagerly():
+ self._resource_deleter = resource_variable_ops.EagerResourceDeleter(
+ handle=self._resource, handle_device="cpu:0")
+ else:
+ ops.add_to_collection(_SUMMARY_WRITER_INIT_COLLECTION_NAME, init_op)
+
+ def init(self):
+ """See `SummaryWriter.init`."""
+ return self._init_op_fn(self._resource)
+
+ def flush(self):
+ """See `SummaryWriter.flush`."""
+ with ops.device("cpu:0"):
+ return gen_summary_ops.flush_summary_writer(self._resource)
+
+ def close(self):
+ """See `SummaryWriter.close`."""
+ with ops.control_dependencies([self.flush()]):
+ with ops.device("cpu:0"):
+ return gen_summary_ops.close_summary_writer(self._resource)
+
+
+class _NoopSummaryWriter(SummaryWriter):
"""A summary writer that does nothing, for create_noop_writer()."""
def set_as_default(self, step=None):
@@ -468,7 +492,7 @@
shared_name = context.shared_name()
else:
shared_name = ops.name_from_scope_name(scope) # pylint: disable=protected-access
- return ResourceSummaryWriter(
+ return _ResourceSummaryWriter(
shared_name=shared_name,
init_op_fn=functools.partial(
gen_summary_ops.create_summary_file_writer,
@@ -476,8 +500,7 @@
max_queue=max_queue,
flush_millis=flush_millis,
filename_suffix=filename_suffix),
- name=name,
- v2=True)
+ name=name)
def create_file_writer(logdir,
@@ -506,7 +529,7 @@
summary writer.
"""
if logdir is None:
- return NoopSummaryWriter()
+ return _NoopSummaryWriter()
logdir = str(logdir)
with ops.device("cpu:0"):
if max_queue is None:
@@ -517,8 +540,9 @@
filename_suffix = constant_op.constant(".v2")
if name is None:
name = "logdir:" + logdir
- return ResourceSummaryWriter(
- shared_name=name,
+ resource = gen_summary_ops.summary_writer(shared_name=name)
+ return _LegacyResourceSummaryWriter(
+ resource=resource,
init_op_fn=functools.partial(
gen_summary_ops.create_summary_file_writer,
logdir=logdir,
@@ -533,7 +557,7 @@
This is useful as a placeholder in code that expects a context manager.
"""
- return NoopSummaryWriter()
+ return _NoopSummaryWriter()
def _cleanse_string(name, pattern, value):
