Upload autosummary.py

dnnlib/tflib/autosummary.py

@@ -0,0 +1,193 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Helper for adding automatically tracked values to Tensorboard.
+
+Autosummary creates an identity op that internally keeps track of the input
+values and automatically shows up in TensorBoard. The reported value
+represents an average over input components. The average is accumulated
+constantly over time and flushed when save_summaries() is called.
+
+Notes:
+- The output tensor must be used as an input for something else in the
+  graph. Otherwise, the autosummary op will not get executed, and the average
+  value will not get accumulated.
+- It is perfectly fine to include autosummaries with the same name in
+  several places throughout the graph, even if they are executed concurrently.
+- It is ok to also pass in a python scalar or numpy array. In this case, it
+  is added to the average immediately.
+"""
+
+from collections import OrderedDict
+import numpy as np
+import tensorflow as tf
+from tensorboard import summary as summary_lib
+from tensorboard.plugins.custom_scalar import layout_pb2
+
+from . import tfutil
+from .tfutil import TfExpression
+from .tfutil import TfExpressionEx
+
+# Enable "Custom scalars" tab in TensorBoard for advanced formatting.
+# Disabled by default to reduce tfevents file size.
+enable_custom_scalars = False
+
+_dtype = tf.float64
+_vars = OrderedDict()  # name => [var, ...]
+_immediate = OrderedDict()  # name => update_op, update_value
+_finalized = False
+_merge_op = None
+
+
+def _create_var(name: str, value_expr: TfExpression) -> TfExpression:
+    """Internal helper for creating autosummary accumulators."""
+    assert not _finalized
+    name_id = name.replace("/", "_")
+    v = tf.cast(value_expr, _dtype)
+
+    if v.shape.is_fully_defined():
+        size = np.prod(v.shape.as_list())
+        size_expr = tf.constant(size, dtype=_dtype)
+    else:
+        size = None
+        size_expr = tf.reduce_prod(tf.cast(tf.shape(v), _dtype))
+
+    if size == 1:
+        if v.shape.ndims != 0:
+            v = tf.reshape(v, [])
+        v = [size_expr, v, tf.square(v)]
+    else:
+        v = [size_expr, tf.reduce_sum(v), tf.reduce_sum(tf.square(v))]
+    v = tf.cond(tf.is_finite(v[1]), lambda: tf.stack(v), lambda: tf.zeros(3, dtype=_dtype))
+
+    with tfutil.absolute_name_scope("Autosummary/" + name_id), tf.control_dependencies(None):
+        var = tf.Variable(tf.zeros(3, dtype=_dtype), trainable=False)  # [sum(1), sum(x), sum(x**2)]
+    update_op = tf.cond(tf.is_variable_initialized(var), lambda: tf.assign_add(var, v), lambda: tf.assign(var, v))
+
+    if name in _vars:
+        _vars[name].append(var)
+    else:
+        _vars[name] = [var]
+    return update_op
+
+
+def autosummary(name: str, value: TfExpressionEx, passthru: TfExpressionEx = None, condition: TfExpressionEx = True) -> TfExpressionEx:
+    """Create a new autosummary.
+
+    Args:
+        name:     Name to use in TensorBoard
+        value:    TensorFlow expression or python value to track
+        passthru: Optionally return this TF node without modifications but tack an autosummary update side-effect to this node.
+
+    Example use of the passthru mechanism:
+
+    n = autosummary('l2loss', loss, passthru=n)
+
+    This is a shorthand for the following code:
+
+    with tf.control_dependencies([autosummary('l2loss', loss)]):
+        n = tf.identity(n)
+    """
+    tfutil.assert_tf_initialized()
+    name_id = name.replace("/", "_")
+
+    if tfutil.is_tf_expression(value):
+        with tf.name_scope("summary_" + name_id), tf.device(value.device):
+            condition = tf.convert_to_tensor(condition, name='condition')
+            update_op = tf.cond(condition, lambda: tf.group(_create_var(name, value)), tf.no_op)
+            with tf.control_dependencies([update_op]):
+                return tf.identity(value if passthru is None else passthru)
+
+    else:  # python scalar or numpy array
+        assert not tfutil.is_tf_expression(passthru)
+        assert not tfutil.is_tf_expression(condition)
+        if condition:
+            if name not in _immediate:
+                with tfutil.absolute_name_scope("Autosummary/" + name_id), tf.device(None), tf.control_dependencies(None):
+                    update_value = tf.placeholder(_dtype)
+                    update_op = _create_var(name, update_value)
+                    _immediate[name] = update_op, update_value
+            update_op, update_value = _immediate[name]
+            tfutil.run(update_op, {update_value: value})
+        return value if passthru is None else passthru
+
+
+def finalize_autosummaries() -> None:
+    """Create the necessary ops to include autosummaries in TensorBoard report.
+    Note: This should be done only once per graph.
+    """
+    global _finalized
+    tfutil.assert_tf_initialized()
+
+    if _finalized:
+        return None
+
+    _finalized = True
+    tfutil.init_uninitialized_vars([var for vars_list in _vars.values() for var in vars_list])
+
+    # Create summary ops.
+    with tf.device(None), tf.control_dependencies(None):
+        for name, vars_list in _vars.items():
+            name_id = name.replace("/", "_")
+            with tfutil.absolute_name_scope("Autosummary/" + name_id):
+                moments = tf.add_n(vars_list)
+                moments /= moments[0]
+                with tf.control_dependencies([moments]):  # read before resetting
+                    reset_ops = [tf.assign(var, tf.zeros(3, dtype=_dtype)) for var in vars_list]
+                    with tf.name_scope(None), tf.control_dependencies(reset_ops):  # reset before reporting
+                        mean = moments[1]
+                        std = tf.sqrt(moments[2] - tf.square(moments[1]))
+                        tf.summary.scalar(name, mean)
+                        if enable_custom_scalars:
+                            tf.summary.scalar("xCustomScalars/" + name + "/margin_lo", mean - std)
+                            tf.summary.scalar("xCustomScalars/" + name + "/margin_hi", mean + std)
+
+    # Setup layout for custom scalars.
+    layout = None
+    if enable_custom_scalars:
+        cat_dict = OrderedDict()
+        for series_name in sorted(_vars.keys()):
+            p = series_name.split("/")
+            cat = p[0] if len(p) >= 2 else ""
+            chart = "/".join(p[1:-1]) if len(p) >= 3 else p[-1]
+            if cat not in cat_dict:
+                cat_dict[cat] = OrderedDict()
+            if chart not in cat_dict[cat]:
+                cat_dict[cat][chart] = []
+            cat_dict[cat][chart].append(series_name)
+        categories = []
+        for cat_name, chart_dict in cat_dict.items():
+            charts = []
+            for chart_name, series_names in chart_dict.items():
+                series = []
+                for series_name in series_names:
+                    series.append(layout_pb2.MarginChartContent.Series(
+                        value=series_name,
+                        lower="xCustomScalars/" + series_name + "/margin_lo",
+                        upper="xCustomScalars/" + series_name + "/margin_hi"))
+                margin = layout_pb2.MarginChartContent(series=series)
+                charts.append(layout_pb2.Chart(title=chart_name, margin=margin))
+            categories.append(layout_pb2.Category(title=cat_name, chart=charts))
+        layout = summary_lib.custom_scalar_pb(layout_pb2.Layout(category=categories))
+    return layout
+
+def save_summaries(file_writer, global_step=None):
+    """Call FileWriter.add_summary() with all summaries in the default graph,
+    automatically finalizing and merging them on the first call.
+    """
+    global _merge_op
+    tfutil.assert_tf_initialized()
+
+    if _merge_op is None:
+        layout = finalize_autosummaries()
+        if layout is not None:
+            file_writer.add_summary(layout)
+        with tf.device(None), tf.control_dependencies(None):
+            _merge_op = tf.summary.merge_all()
+
+    file_writer.add_summary(_merge_op.eval(), global_step)