Utility functions for tensorflow and tensorboard

tensorflow_util.py

import tensorflow as tf
from tensorflow.core.framework import summary_pb2


pprint = lambda x, msg: tf.Print(x, [x], message=msg)


def scalar_summary(tag, value):
  """Create a scalar summary and add it to the collection if it doesn't exist."""
  if summary_exists(tag) is None:
    summ = tf.scalar_summary(tag, value)
    tf.add_to_collection("summary_tags", {"tag": tag, "tensor": summ})


def summary_exists(tag):
    """ summary_exists.
    Check if a summary exists.
    Arguments:
        tag: `str`. The summary name.
    Returns:
        A `bool`. Whether the summary exists or not.
    FROM https://github.com/tflearn/tflearn/blob/master/tflearn/summaries.py
    """
    return next(
        (item['tensor'] for item in tf.get_collection("summary_tags") if
         item["tag"] == tag), None)
  


def get_pos_neg_means(tensor):
  """Return the means of the positive and negative parts of a tensor."""
  # flatten the tensor
  t = tf.reshape(tensor, [-1])
  # sort the tensor. segment_mean needs sorted indices
  sorted_t, _ = tf.nn.top_k(t, t.get_shape()[0])
  sorted_t = tf.reverse(sorted_t, [True])
  mean = tf.reduce_mean(sorted_t)
  means = tf.cond(tf.equal(mean, 1.),
      lambda: tf.reduce_mean(t),
      lambda: tf.segment_mean(sorted_t, tf.cast(sorted_t > 0, tf.int32)))
  pos_mean = tf.reduce_sum(tf.maximum(means, 0.))
  neg_mean = tf.reduce_sum(tf.minimum(means, 0.))
  return pos_mean, neg_mean

def make_summary(tag, value):
  """Return a manual protobuf summary using numpy arrays for tensorboard."""
  return summary_pb2.Summary(value=[summary_pb2.Summary.Value(
      tag=tag, simple_value=value)])


def make_and_add_summaries(summary_writer, tags_and_values, step):
  """Make multiple summaries and add them to the writer."""
  for tag_value in tags_and_values:
    tag, value = tag_value
    summary_str = make_summary(tag, float(value))
    summary_writer.add_summary(summary_str, global_step=step)
    
def get_grad_variance(sess, grad, feed_dict):
  """Returns the variance of the gradients."""
  np_grad_list = [] * len(grad)
  for sample in range(10):
    np_grad_list.append(sess.run(grad, feed_dict))
  np_grads_per_var = [[g[v] for g in np_grad_list] for v in range(len(grad))]
  np_grads_per_var_stacked = [np.vstack(g_list) for g_list in np_grads_per_var]
  # get the variance across samples
  variances = [np.mean(np.var(g_var, axis=0)) for g_var in np_grads_per_var_stacked]
  # return the mean variance
  return np.mean(variances)


def get_mean_tensor(sess, tensor, feed_dict, n_samples=100):
  """Get mean of a tensor."""
  t_lst = []
  for sample in range(n_samples):
    t_lst.append(sess.run(tensor, feed_dict))
  return np.mean(np.vstack(t_lst), axis=0)