hparams: write hparams_config summary directly (#2139)

Summary:
The `Experiment` object bundled `HParam`s and `Metric`s with some
metadata that’s not actually used in the current UI. We don’t think that
it pulls its conceptual weight, so this commit replaces it with a direct
summary-writing operation.

This function will soon be extracted into a `summary_pb2` module, as
part of a [larger plan to refactor the `api` module][1]. Making this
change first minimizes churn in the demo code.

[1]: #2139 (comment)

Cf. #1998.

Test Plan:
Unit tests modified appropriately, and the demo still works.

wchargin-branch: hparams-experiment-writing

tensorboard/plugins/hparams/api.py

@@ -31,90 +31,44 @@
 from tensorboard.plugins.hparams import summary
 
 
-class Experiment(object):
-  """A top-level experiment description.
+def hparams_config(hparams, metrics, time_created_secs=None):
+  """Write a top-level experiment configuration.
 
-  An experiment has a fixed set of hyperparameters and metrics, and
-  consists of multiple sessions. Each session has different associated
-  hyperparameter values and metric values.
-  """
-
-  def __init__(
-      self,
-      hparams,
-      metrics,
-      user=None,
-      description=None,
-      time_created_secs=None,
-  ):
-    """Create an experiment object.
-
-    Args:
-      hparams: A list of `HParam` values.
-      metrics: A list of `Metric` values.
-      user: An optional string denoting the user or group that owns this
-        experiment.
-      description: An optional Markdown string describing this
-        experiment.
-      time_created_secs: The time that this experiment was created, as
-        seconds since epoch. Defaults to the current time.
-    """
-    self._hparams = list(hparams)
-    self._metrics = list(metrics)
-    self._user = user
-    self._description = description
-    if time_created_secs is None:
-      time_created_secs = time.time()
-    self._time_created_secs = time_created_secs
-
-  @property
-  def hparams(self):
-    return list(self._hparams)
-
-  @property
-  def metrics(self):
-    return list(self._metrics)
+  This configuration describes the hyperparameters and metrics that will
+  be tracked in the experiment, but does not record any actual values of
+  those hyperparameters and metrics. It can be created before any models
+  are actually trained.
 
-  @property
-  def user(self):
-    return self._user
-
-  @property
-  def description(self):
-    return self._description
-
-  @property
-  def time_created_secs(self):
-    return self._time_created_secs
-
-  def summary_pb(self):
-    """Create a top-level experiment summary describing this experiment.
-
-    The resulting summary should be written to a log directory that
-    encloses all the individual sessions' log directories.
-
-    Analogous to the low-level `experiment_pb` function in the
-    `hparams.summary` module.
-    """
-    hparam_infos = []
-    for hparam in self._hparams:
-      info = api_pb2.HParamInfo(
-          name=hparam.name,
-          description=hparam.description,
-          display_name=hparam.display_name,
-      )
-      domain = hparam.domain
-      if domain is not None:
-        domain.update_hparam_info(info)
-      hparam_infos.append(info)
-    metric_infos = [metric.as_proto() for metric in self._metrics]
-    return summary.experiment_pb(
-        hparam_infos=hparam_infos,
-        metric_infos=metric_infos,
-        user=self._user,
-        description=self._description,
-        time_created_secs=self._time_created_secs,
+  Args:
+    hparams: A list of `HParam` values.
+    metrics: A list of `Metric` values.
+    time_created_secs: The time that this experiment was created, as
+      seconds since epoch. Defaults to the current time.
+  """
+  hparam_infos = []
+  for hparam in hparams:
+    info = api_pb2.HParamInfo(
+        name=hparam.name,
+        description=hparam.description,
+        display_name=hparam.display_name,
     )
+    domain = hparam.domain
+    if domain is not None:
+      domain.update_hparam_info(info)
+    hparam_infos.append(info)
+  metric_infos = [metric.as_proto() for metric in metrics]
+  experiment_pb = summary.experiment_pb(
+      hparam_infos=hparam_infos,
+      metric_infos=metric_infos,
+      time_created_secs=time_created_secs,
+  )
+  raw_pb = experiment_pb.SerializeToString()
+  summary_scope = (
+      getattr(tf.compat.v2.summary.experimental, "summary_scope", None)
+      or tf.summary.summary_scope
+  )
+  with summary_scope("hparams_summary"):
+    return tf.compat.v2.summary.experimental.write_raw_pb(raw_pb, step=0)
 
 
 class HParam(object):
@@ -439,7 +393,7 @@ def __init__(
       logdir: The log directory for this session.
       hparams: A `dict` mapping hyperparameters to the values used in
         this session. Keys should be the names of `HParam` objects used
-        in an `Experiment`, or the `HParam` objects themselves. Values
+        in an experiment, or the `HParam` objects themselves. Values
         should be Python `bool`, `int`, `float`, or `string` values,
         depending on the type of the hyperparameter.
       group_name: The name of the session group containing this session,

tensorboard/plugins/hparams/api_test.py

@@ -42,9 +42,11 @@
 tf.compat.v1.enable_eager_execution()
 
 
-class ExperimentTest(test.TestCase):
-  def test_summary_pb(self):
-    hparams = [
+class HParamsConfigTest(test.TestCase):
+  def setUp(self):
+    self.logdir = os.path.join(self.get_temp_dir(), "logs")
+
+    self.hparams = [
         hp.HParam("learning_rate", hp.RealInterval(1e-2, 1e-1)),
         hp.HParam("dense_layers", hp.IntInterval(2, 7)),
         hp.HParam("optimizer", hp.Discrete(["adam", "sgd"])),
@@ -56,7 +58,7 @@ def test_summary_pb(self):
             description="descriptive",
         ),
     ]
-    metrics = [
+    self.metrics = [
         hp.Metric("samples_per_second"),
         hp.Metric(group="train", tag="batch_loss", display_name="loss (train)"),
         hp.Metric(
@@ -67,25 +69,11 @@ def test_summary_pb(self):
             dataset_type=hp.Metric.VALIDATION,
         ),
     ]
-    experiment = hp.Experiment(
-        hparams=hparams,
-        metrics=metrics,
-        user="zalgo",
-        description="nothing to see here; move along",
-        time_created_secs=1555624767,
-    )
-
-    self.assertEqual(experiment.hparams, hparams)
-    self.assertEqual(experiment.metrics, metrics)
-    self.assertEqual(experiment.user, "zalgo"),
-    self.assertEqual(experiment.description, "nothing to see here; move along")
-    self.assertEqual(experiment.time_created_secs, 1555624767)
+    self.time_created_secs = 1555624767.0
 
-    expected_experiment_pb = api_pb2.Experiment()
+    self.expected_experiment_pb = api_pb2.Experiment()
     text_format.Merge(
         """
-        description: "nothing to see here; move along"
-        user: "zalgo"
         time_created_secs: 1555624767.0
         hparam_infos {
           name: "learning_rate"
@@ -154,14 +142,80 @@ def test_summary_pb(self):
           dataset_type: DATASET_VALIDATION
         }
         """,
-        expected_experiment_pb,
+        self.expected_experiment_pb,
+    )
+
+  def _get_unique_summary_value(self, logdir):
+    """Get the unique summary `Value` stored in `logdir`.
+
+    Specifically, `logdir` must be a directory containing exactly one
+    entry, which must be an events file of whose events exactly one is a
+    summary, which must have exactly one `value`. This unique `value`
+    will be returned.
+
+    Args:
+      logdir: String path to a logdir.
+
+    Returns:
+      A `summary_pb2.Summary.Value` object.
+    """
+    files = os.listdir(logdir)
+    self.assertEqual(len(files), 1, files)
+    events_file = os.path.join(logdir, files[0])
+    summaries = [
+        event.summary
+        for event in tf.compat.v1.train.summary_iterator(events_file)
+        if event.WhichOneof("what") == "summary"
+    ]
+    self.assertEqual(len(summaries), 1, summaries)
+    values = summaries[0].value
+    self.assertEqual(len(values), 1, values)
+    return values[0]
+
+  def _check_logdir(self, logdir):
+    """Test that the experiment summary was written to `logdir`."""
+    actual_value = self._get_unique_summary_value(logdir)
+    self.assertEqual(
+        actual_value.metadata.plugin_data.plugin_name,
+        metadata.PLUGIN_NAME,
     )
-    actual_summary_pb = experiment.summary_pb()
-    plugin_content = actual_summary_pb.value[0].metadata.plugin_data.content
+    plugin_content = actual_value.metadata.plugin_data.content
     self.assertEqual(
         metadata.parse_experiment_plugin_data(plugin_content),
-        expected_experiment_pb,
+        self.expected_experiment_pb,
+    )
+
+  def test_eager(self):
+    with tf.compat.v2.summary.create_file_writer(self.logdir).as_default():
+      result = hp.hparams_config(
+          hparams=self.hparams,
+          metrics=self.metrics,
+          time_created_secs=self.time_created_secs,
+      )
+      self.assertTrue(result)
+    self._check_logdir(self.logdir)
+
+  def test_graph_mode(self):
+    with \
+        tf.compat.v1.Graph().as_default(), \
+        tf.compat.v1.Session() as sess, \
+        tf.compat.v2.summary.create_file_writer(self.logdir).as_default() as w:
+      sess.run(w.init())
+      summ = hp.hparams_config(
+          hparams=self.hparams,
+          metrics=self.metrics,
+          time_created_secs=self.time_created_secs,
+      )
+      self.assertTrue(sess.run(summ))
+    self._check_logdir(self.logdir)
+
+  def test_eager_no_default_writer(self):
+    result = hp.hparams_config(
+        hparams=self.hparams,
+        metrics=self.metrics,
+        time_created_secs=self.time_created_secs,
     )
+    self.assertFalse(result)  # no default writer
 
 
 class IntIntervalTest(test.TestCase):

tensorboard/plugins/hparams/hparams_demo.py

@@ -215,13 +215,8 @@ def run_all(logdir, verbose=False):
   data = prepare_data()
   rng = random.Random(0)
 
-  base_writer = tf.summary.create_file_writer(logdir)
-  with base_writer.as_default():
-    experiment = hp.Experiment(hparams=HPARAMS, metrics=METRICS)
-    experiment_string = experiment.summary_pb().SerializeToString()
-    tf.summary.experimental.write_raw_pb(experiment_string, step=0)
-    base_writer.flush()
-  base_writer.close()
+  with tf.summary.create_file_writer(logdir).as_default():
+    hp.hparams_config(hparams=HPARAMS, metrics=METRICS)
 
   sessions_per_group = 2
   num_sessions = flags.FLAGS.num_session_groups * sessions_per_group