|
| 1 | +# Copyright (c) SenseTime Research. All rights reserved. |
| 2 | + |
| 3 | +# Copyright (c) 2019, NVIDIA Corporation. All rights reserved. |
| 4 | +# |
| 5 | +# This work is made available under the Nvidia Source Code License-NC. |
| 6 | +# To view a copy of this license, visit |
| 7 | +# https://nvlabs.github.io/stylegan2/license.html |
| 8 | + |
| 9 | +"""Helper for adding automatically tracked values to Tensorboard. |
| 10 | +
|
| 11 | +Autosummary creates an identity op that internally keeps track of the input |
| 12 | +values and automatically shows up in TensorBoard. The reported value |
| 13 | +represents an average over input components. The average is accumulated |
| 14 | +constantly over time and flushed when save_summaries() is called. |
| 15 | +
|
| 16 | +Notes: |
| 17 | +- The output tensor must be used as an input for something else in the |
| 18 | + graph. Otherwise, the autosummary op will not get executed, and the average |
| 19 | + value will not get accumulated. |
| 20 | +- It is perfectly fine to include autosummaries with the same name in |
| 21 | + several places throughout the graph, even if they are executed concurrently. |
| 22 | +- It is ok to also pass in a python scalar or numpy array. In this case, it |
| 23 | + is added to the average immediately. |
| 24 | +""" |
| 25 | + |
| 26 | +from collections import OrderedDict |
| 27 | +import numpy as np |
| 28 | +import tensorflow as tf |
| 29 | +from tensorboard import summary as summary_lib |
| 30 | +from tensorboard.plugins.custom_scalar import layout_pb2 |
| 31 | + |
| 32 | +from . import tfutil |
| 33 | +from .tfutil import TfExpression |
| 34 | +from .tfutil import TfExpressionEx |
| 35 | + |
| 36 | +# Enable "Custom scalars" tab in TensorBoard for advanced formatting. |
| 37 | +# Disabled by default to reduce tfevents file size. |
| 38 | +enable_custom_scalars = False |
| 39 | + |
| 40 | +_dtype = tf.float64 |
| 41 | +_vars = OrderedDict() # name => [var, ...] |
| 42 | +_immediate = OrderedDict() # name => update_op, update_value |
| 43 | +_finalized = False |
| 44 | +_merge_op = None |
| 45 | + |
| 46 | + |
| 47 | +def _create_var(name: str, value_expr: TfExpression) -> TfExpression: |
| 48 | + """Internal helper for creating autosummary accumulators.""" |
| 49 | + assert not _finalized |
| 50 | + name_id = name.replace("/", "_") |
| 51 | + v = tf.cast(value_expr, _dtype) |
| 52 | + |
| 53 | + if v.shape.is_fully_defined(): |
| 54 | + size = np.prod(v.shape.as_list()) |
| 55 | + size_expr = tf.constant(size, dtype=_dtype) |
| 56 | + else: |
| 57 | + size = None |
| 58 | + size_expr = tf.reduce_prod(tf.cast(tf.shape(v), _dtype)) |
| 59 | + |
| 60 | + if size == 1: |
| 61 | + if v.shape.ndims != 0: |
| 62 | + v = tf.reshape(v, []) |
| 63 | + v = [size_expr, v, tf.square(v)] |
| 64 | + else: |
| 65 | + v = [size_expr, tf.reduce_sum(v), tf.reduce_sum(tf.square(v))] |
| 66 | + v = tf.cond(tf.is_finite(v[1]), lambda: tf.stack(v), lambda: tf.zeros(3, dtype=_dtype)) |
| 67 | + |
| 68 | + with tfutil.absolute_name_scope("Autosummary/" + name_id), tf.control_dependencies(None): |
| 69 | + var = tf.Variable(tf.zeros(3, dtype=_dtype), trainable=False) # [sum(1), sum(x), sum(x**2)] |
| 70 | + update_op = tf.cond(tf.is_variable_initialized(var), lambda: tf.assign_add(var, v), lambda: tf.assign(var, v)) |
| 71 | + |
| 72 | + if name in _vars: |
| 73 | + _vars[name].append(var) |
| 74 | + else: |
| 75 | + _vars[name] = [var] |
| 76 | + return update_op |
| 77 | + |
| 78 | + |
| 79 | +def autosummary(name: str, value: TfExpressionEx, passthru: TfExpressionEx = None, condition: TfExpressionEx = True) -> TfExpressionEx: |
| 80 | + """Create a new autosummary. |
| 81 | +
|
| 82 | + Args: |
| 83 | + name: Name to use in TensorBoard |
| 84 | + value: TensorFlow expression or python value to track |
| 85 | + passthru: Optionally return this TF node without modifications but tack an autosummary update side-effect to this node. |
| 86 | +
|
| 87 | + Example use of the passthru mechanism: |
| 88 | +
|
| 89 | + n = autosummary('l2loss', loss, passthru=n) |
| 90 | +
|
| 91 | + This is a shorthand for the following code: |
| 92 | +
|
| 93 | + with tf.control_dependencies([autosummary('l2loss', loss)]): |
| 94 | + n = tf.identity(n) |
| 95 | + """ |
| 96 | + tfutil.assert_tf_initialized() |
| 97 | + name_id = name.replace("/", "_") |
| 98 | + |
| 99 | + if tfutil.is_tf_expression(value): |
| 100 | + with tf.name_scope("summary_" + name_id), tf.device(value.device): |
| 101 | + condition = tf.convert_to_tensor(condition, name='condition') |
| 102 | + update_op = tf.cond(condition, lambda: tf.group(_create_var(name, value)), tf.no_op) |
| 103 | + with tf.control_dependencies([update_op]): |
| 104 | + return tf.identity(value if passthru is None else passthru) |
| 105 | + |
| 106 | + else: # python scalar or numpy array |
| 107 | + assert not tfutil.is_tf_expression(passthru) |
| 108 | + assert not tfutil.is_tf_expression(condition) |
| 109 | + if condition: |
| 110 | + if name not in _immediate: |
| 111 | + with tfutil.absolute_name_scope("Autosummary/" + name_id), tf.device(None), tf.control_dependencies(None): |
| 112 | + update_value = tf.placeholder(_dtype) |
| 113 | + update_op = _create_var(name, update_value) |
| 114 | + _immediate[name] = update_op, update_value |
| 115 | + update_op, update_value = _immediate[name] |
| 116 | + tfutil.run(update_op, {update_value: value}) |
| 117 | + return value if passthru is None else passthru |
| 118 | + |
| 119 | + |
| 120 | +def finalize_autosummaries() -> None: |
| 121 | + """Create the necessary ops to include autosummaries in TensorBoard report. |
| 122 | + Note: This should be done only once per graph. |
| 123 | + """ |
| 124 | + global _finalized |
| 125 | + tfutil.assert_tf_initialized() |
| 126 | + |
| 127 | + if _finalized: |
| 128 | + return None |
| 129 | + |
| 130 | + _finalized = True |
| 131 | + tfutil.init_uninitialized_vars([var for vars_list in _vars.values() for var in vars_list]) |
| 132 | + |
| 133 | + # Create summary ops. |
| 134 | + with tf.device(None), tf.control_dependencies(None): |
| 135 | + for name, vars_list in _vars.items(): |
| 136 | + name_id = name.replace("/", "_") |
| 137 | + with tfutil.absolute_name_scope("Autosummary/" + name_id): |
| 138 | + moments = tf.add_n(vars_list) |
| 139 | + moments /= moments[0] |
| 140 | + with tf.control_dependencies([moments]): # read before resetting |
| 141 | + reset_ops = [tf.assign(var, tf.zeros(3, dtype=_dtype)) for var in vars_list] |
| 142 | + with tf.name_scope(None), tf.control_dependencies(reset_ops): # reset before reporting |
| 143 | + mean = moments[1] |
| 144 | + std = tf.sqrt(moments[2] - tf.square(moments[1])) |
| 145 | + tf.summary.scalar(name, mean) |
| 146 | + if enable_custom_scalars: |
| 147 | + tf.summary.scalar("xCustomScalars/" + name + "/margin_lo", mean - std) |
| 148 | + tf.summary.scalar("xCustomScalars/" + name + "/margin_hi", mean + std) |
| 149 | + |
| 150 | + # Setup layout for custom scalars. |
| 151 | + layout = None |
| 152 | + if enable_custom_scalars: |
| 153 | + cat_dict = OrderedDict() |
| 154 | + for series_name in sorted(_vars.keys()): |
| 155 | + p = series_name.split("/") |
| 156 | + cat = p[0] if len(p) >= 2 else "" |
| 157 | + chart = "/".join(p[1:-1]) if len(p) >= 3 else p[-1] |
| 158 | + if cat not in cat_dict: |
| 159 | + cat_dict[cat] = OrderedDict() |
| 160 | + if chart not in cat_dict[cat]: |
| 161 | + cat_dict[cat][chart] = [] |
| 162 | + cat_dict[cat][chart].append(series_name) |
| 163 | + categories = [] |
| 164 | + for cat_name, chart_dict in cat_dict.items(): |
| 165 | + charts = [] |
| 166 | + for chart_name, series_names in chart_dict.items(): |
| 167 | + series = [] |
| 168 | + for series_name in series_names: |
| 169 | + series.append(layout_pb2.MarginChartContent.Series( |
| 170 | + value=series_name, |
| 171 | + lower="xCustomScalars/" + series_name + "/margin_lo", |
| 172 | + upper="xCustomScalars/" + series_name + "/margin_hi")) |
| 173 | + margin = layout_pb2.MarginChartContent(series=series) |
| 174 | + charts.append(layout_pb2.Chart(title=chart_name, margin=margin)) |
| 175 | + categories.append(layout_pb2.Category(title=cat_name, chart=charts)) |
| 176 | + layout = summary_lib.custom_scalar_pb(layout_pb2.Layout(category=categories)) |
| 177 | + return layout |
| 178 | + |
| 179 | +def save_summaries(file_writer, global_step=None): |
| 180 | + """Call FileWriter.add_summary() with all summaries in the default graph, |
| 181 | + automatically finalizing and merging them on the first call. |
| 182 | + """ |
| 183 | + global _merge_op |
| 184 | + tfutil.assert_tf_initialized() |
| 185 | + |
| 186 | + if _merge_op is None: |
| 187 | + layout = finalize_autosummaries() |
| 188 | + if layout is not None: |
| 189 | + file_writer.add_summary(layout) |
| 190 | + with tf.device(None), tf.control_dependencies(None): |
| 191 | + _merge_op = tf.summary.merge_all() |
| 192 | + |
| 193 | + file_writer.add_summary(_merge_op.eval(), global_step) |
0 commit comments