[tune] TensorBoard HParams for TF2.0 (#5678)

python/ray/tune/logger.py

@@ -38,9 +38,10 @@ class Logger(object):
         logdir: Directory for all logger creators to log to.
     """
 
-    def __init__(self, config, logdir):
+    def __init__(self, config, logdir, trial=None):
         self.config = config
         self.logdir = logdir
+        self.trial = trial
         self._init()
 
     def _init(self):
@@ -135,7 +136,7 @@ class JsonLogger(Logger):
             cloudpickle.dump(self.config, f)
 
 
-def tf2_compat_logger(config, logdir):
+def tf2_compat_logger(config, logdir, trial=None):
     """Chooses TensorBoard logger depending on imported TF version."""
     global tf
     if "RLLIB_TEST_NO_TF_IMPORT" in os.environ:
@@ -148,9 +149,9 @@ def tf2_compat_logger(config, logdir):
                        distutils.version.LooseVersion("2.0.0"))
         if use_tf2_api:
             tf = tf.compat.v2  # setting this for TF2.0
-            return TF2Logger(config, logdir)
+            return TF2Logger(config, logdir, trial)
         else:
-            return TFLogger(config, logdir)
+            return TFLogger(config, logdir, trial)
 
 
 class TF2Logger(Logger):
@@ -166,10 +167,12 @@ class TF2Logger(Logger):
 
     def _init(self):
         self._file_writer = None
+        self._hp_logged = False
 
     def on_result(self, result):
         if self._file_writer is None:
             from tensorflow.python.eager import context
+            from tensorboard.plugins.hparams import api as hp
             self._context = context
             self._file_writer = tf.summary.create_file_writer(self.logdir)
         with tf.device("/CPU:0"), self._context.eager_mode():
@@ -178,6 +181,16 @@ class TF2Logger(Logger):
                     TIMESTEPS_TOTAL) or result[TRAINING_ITERATION]
 
                 tmp = result.copy()
+                if not self._hp_logged:
+                    if self.trial and self.trial.evaluated_params:
+                        try:
+                            hp.hparams(
+                                self.trial.evaluated_params,
+                                trial_id=self.trial.trial_id)
+                        except Exception as exc:
+                            logger.error("HParams failed with %s", exc)
+                    self._hp_logged = True
+
                 for k in [
                         "config", "pid", "timestamp", TIME_TOTAL_S,
                         TRAINING_ITERATION
@@ -305,7 +318,12 @@ class UnifiedLogger(Logger):
             See ray/python/ray/tune/log_sync.py
     """
 
-    def __init__(self, config, logdir, loggers=None, sync_function=None):
+    def __init__(self,
+                 config,
+                 logdir,
+                 trial=None,
+                 loggers=None,
+                 sync_function=None):
         if loggers is None:
             self._logger_cls_list = DEFAULT_LOGGERS
         else:
@@ -313,13 +331,13 @@ class UnifiedLogger(Logger):
         self._sync_function = sync_function
         self._log_syncer = None
 
-        super(UnifiedLogger, self).__init__(config, logdir)
+        super(UnifiedLogger, self).__init__(config, logdir, trial)
 
     def _init(self):
         self._loggers = []
         for cls in self._logger_cls_list:
             try:
-                self._loggers.append(cls(self.config, self.logdir))
+                self._loggers.append(cls(self.config, self.logdir, self.trial))
             except Exception as exc:
                 logger.warning("Could not instantiate {}: {}.".format(
                     cls.__name__, str(exc)))

python/ray/tune/suggest/basic_variant.py

@@ -8,7 +8,8 @@ import random
 from ray.tune.error import TuneError
 from ray.tune.experiment import convert_to_experiment_list
 from ray.tune.config_parser import make_parser, create_trial_from_spec
-from ray.tune.suggest.variant_generator import generate_variants
+from ray.tune.suggest.variant_generator import (generate_variants, format_vars,
+                                                flatten_resolved_vars)
 from ray.tune.suggest.search import SearchAlgorithm
 
 
@@ -78,13 +79,13 @@ class BasicVariantGenerator(SearchAlgorithm):
             for resolved_vars, spec in generate_variants(unresolved_spec):
                 experiment_tag = str(self._counter)
                 if resolved_vars:
-                    experiment_tag += "_{}".format(resolved_vars)
+                    experiment_tag += "_{}".format(format_vars(resolved_vars))
                 self._counter += 1
                 yield create_trial_from_spec(
                     spec,
                     output_path,
                     self._parser,
-                    evaluated_params=resolved_vars,
+                    evaluated_params=flatten_resolved_vars(resolved_vars),
                     experiment_tag=experiment_tag)
 
     def is_finished(self):

python/ray/tune/suggest/suggestion.py

@@ -7,7 +7,7 @@ import copy
 
 from ray.tune.error import TuneError
 from ray.tune.trial import Trial
-from ray.tune.util import merge_dicts
+from ray.tune.util import merge_dicts, flatten_dict
 from ray.tune.experiment import convert_to_experiment_list
 from ray.tune.config_parser import make_parser, create_trial_from_spec
 from ray.tune.suggest.search import SearchAlgorithm
@@ -89,7 +89,8 @@ class SuggestionAlgorithm(SearchAlgorithm):
                 else:
                     break
             spec = copy.deepcopy(experiment_spec)
-            spec["config"] = merge_dicts(spec["config"], suggested_config)
+            spec["config"] = merge_dicts(spec["config"],
+                                         copy.deepcopy(suggested_config))
             flattened_config = resolve_nested_dict(spec["config"])
             self._counter += 1
             tag = "{0}_{1}".format(
@@ -98,7 +99,7 @@ class SuggestionAlgorithm(SearchAlgorithm):
                 spec,
                 output_path,
                 self._parser,
-                evaluated_params=list(suggested_config),
+                evaluated_params=flatten_dict(suggested_config),
                 experiment_tag=tag,
                 trial_id=trial_id)
 

python/ray/tune/suggest/variant_generator.py

@@ -39,10 +39,15 @@ def generate_variants(unresolved_spec):
 
         "activation": {"grid_search": ["relu", "tanh"]}
         "cpu": {"eval": "spec.config.num_workers"}
+
+    Use `format_vars` to format the returned dict of hyperparameters.
+
+    Yields:
+        (Dict of resolved variables, Spec object)
     """
     for resolved_vars, spec in _generate_variants(unresolved_spec):
         assert not _unresolved_values(spec)
-        yield format_vars(resolved_vars), spec
+        yield resolved_vars, spec
 
 
 def grid_search(values):
@@ -79,6 +84,7 @@ def resolve_nested_dict(nested_dict):
 
 
 def format_vars(resolved_vars):
+    """Formats the resolved variable dict into a single string."""
     out = []
     for path, value in sorted(resolved_vars.items()):
         if path[0] in ["run", "env", "resources_per_trial"]:
@@ -96,6 +102,17 @@ def format_vars(resolved_vars):
     return ",".join(out)
 
 
+def flatten_resolved_vars(resolved_vars):
+    """Formats the resolved variable dict into a mapping of (str -> value)."""
+    flattened_resolved_vars_dict = {}
+    for pieces, value in resolved_vars.items():
+        if pieces[0] == "config":
+            pieces = pieces[1:]
+        pieces = [str(piece) for piece in pieces]
+        flattened_resolved_vars_dict["/".join(pieces)] = value
+    return flattened_resolved_vars_dict
+
+
 def _clean_value(value):
     if isinstance(value, float):
         return "{:.5}".format(value)

python/ray/tune/tests/test_trial_runner.py

@@ -1202,6 +1202,7 @@ class VariantGeneratorTest(unittest.TestCase):
         self.assertEqual(trials[0].trainable_name, "PPO")
         self.assertEqual(trials[0].experiment_tag, "0")
         self.assertEqual(trials[0].max_failures, 5)
+        self.assertEqual(trials[0].evaluated_params, {})
         self.assertEqual(trials[0].local_dir,
                          os.path.join(DEFAULT_RESULTS_DIR, "tune-pong"))
         self.assertEqual(trials[1].experiment_tag, "1")
@@ -1218,6 +1219,7 @@ class VariantGeneratorTest(unittest.TestCase):
         trials = list(trials)
         self.assertEqual(len(trials), 1)
         self.assertEqual(trials[0].config, {"foo": 4})
+        self.assertEqual(trials[0].evaluated_params, {"foo": 4})
         self.assertEqual(trials[0].experiment_tag, "0_foo=4")
 
     def testGridSearch(self):
@@ -1230,18 +1232,72 @@ class VariantGeneratorTest(unittest.TestCase):
                 "foo": {
                     "grid_search": [1, 2, 3]
                 },
+                "baz": "asd",
             },
         }, "grid_search")
         trials = list(trials)
         self.assertEqual(len(trials), 6)
-        self.assertEqual(trials[0].config, {"bar": True, "foo": 1})
+        self.assertEqual(trials[0].config, {
+            "bar": True,
+            "foo": 1,
+            "baz": "asd",
+        })
+        self.assertEqual(trials[0].evaluated_params, {
+            "bar": True,
+            "foo": 1,
+        })
         self.assertEqual(trials[0].experiment_tag, "0_bar=True,foo=1")
-        self.assertEqual(trials[1].config, {"bar": False, "foo": 1})
+
+        self.assertEqual(trials[1].config, {
+            "bar": False,
+            "foo": 1,
+            "baz": "asd",
+        })
+        self.assertEqual(trials[1].evaluated_params, {
+            "bar": False,
+            "foo": 1,
+        })
         self.assertEqual(trials[1].experiment_tag, "1_bar=False,foo=1")
-        self.assertEqual(trials[2].config, {"bar": True, "foo": 2})
-        self.assertEqual(trials[3].config, {"bar": False, "foo": 2})
-        self.assertEqual(trials[4].config, {"bar": True, "foo": 3})
-        self.assertEqual(trials[5].config, {"bar": False, "foo": 3})
+
+        self.assertEqual(trials[2].config, {
+            "bar": True,
+            "foo": 2,
+            "baz": "asd",
+        })
+        self.assertEqual(trials[2].evaluated_params, {
+            "bar": True,
+            "foo": 2,
+        })
+
+        self.assertEqual(trials[3].config, {
+            "bar": False,
+            "foo": 2,
+            "baz": "asd",
+        })
+        self.assertEqual(trials[3].evaluated_params, {
+            "bar": False,
+            "foo": 2,
+        })
+
+        self.assertEqual(trials[4].config, {
+            "bar": True,
+            "foo": 3,
+            "baz": "asd",
+        })
+        self.assertEqual(trials[4].evaluated_params, {
+            "bar": True,
+            "foo": 3,
+        })
+
+        self.assertEqual(trials[5].config, {
+            "bar": False,
+            "foo": 3,
+            "baz": "asd",
+        })
+        self.assertEqual(trials[5].evaluated_params, {
+            "bar": False,
+            "foo": 3,
+        })
 
     def testGridSearchAndEval(self):
         trials = self.generate_trials({
@@ -1250,11 +1306,22 @@ class VariantGeneratorTest(unittest.TestCase):
                 "qux": tune.sample_from(lambda spec: 2 + 2),
                 "bar": grid_search([True, False]),
                 "foo": grid_search([1, 2, 3]),
+                "baz": "asd",
             },
         }, "grid_eval")
         trials = list(trials)
         self.assertEqual(len(trials), 6)
-        self.assertEqual(trials[0].config, {"bar": True, "foo": 1, "qux": 4})
+        self.assertEqual(trials[0].config, {
+            "bar": True,
+            "foo": 1,
+            "qux": 4,
+            "baz": "asd",
+        })
+        self.assertEqual(trials[0].evaluated_params, {
+            "bar": True,
+            "foo": 1,
+            "qux": 4,
+        })
         self.assertEqual(trials[0].experiment_tag, "0_bar=True,foo=1,qux=4")
 
     def testConditionResolution(self):
@@ -1269,6 +1336,8 @@ class VariantGeneratorTest(unittest.TestCase):
         trials = list(trials)
         self.assertEqual(len(trials), 1)
         self.assertEqual(trials[0].config, {"x": 1, "y": 2, "z": 3})
+        self.assertEqual(trials[0].evaluated_params, {"y": 2, "z": 3})
+        self.assertEqual(trials[0].experiment_tag, "0_y=2,z=3")
 
     def testDependentLambda(self):
         trials = self.generate_trials({
@@ -1299,6 +1368,36 @@ class VariantGeneratorTest(unittest.TestCase):
         self.assertEqual(trials[0].config, {"x": 100, "y": 1})
         self.assertEqual(trials[1].config, {"x": 200, "y": 1})
 
+    def testNestedValues(self):
+        trials = self.generate_trials({
+            "run": "PPO",
+            "config": {
+                "x": {
+                    "y": {
+                        "z": tune.sample_from(lambda spec: 1)
+                    }
+                },
+                "y": tune.sample_from(lambda spec: 12),
+                "z": tune.sample_from(lambda spec: spec.config.x.y.z * 100),
+            },
+        }, "nested_values")
+        trials = list(trials)
+        self.assertEqual(len(trials), 1)
+        self.assertEqual(trials[0].config, {
+            "x": {
+                "y": {
+                    "z": 1
+                }
+            },
+            "y": 12,
+            "z": 100
+        })
+        self.assertEqual(trials[0].evaluated_params, {
+            "x/y/z": 1,
+            "y": 12,
+            "z": 100
+        })
+
     def testLogUniform(self):
         sampler = tune.loguniform(1e-10, 1e-1).func
         results = [sampler(None) for i in range(1000)]

python/ray/tune/trainable.py

@@ -81,8 +81,8 @@ class Trainable(object):
                 os.makedirs(DEFAULT_RESULTS_DIR)
             self._logdir = tempfile.mkdtemp(
                 prefix=logdir_prefix, dir=DEFAULT_RESULTS_DIR)
-            self._result_logger = UnifiedLogger(self.config, self._logdir,
-                                                None)
+            self._result_logger = UnifiedLogger(
+                self.config, self._logdir, loggers=None)
 
         self._iteration = 0
         self._time_total = 0.0

python/ray/tune/trial.py

@@ -136,7 +136,7 @@ class Trial(object):
         self.local_dir = local_dir  # This remains unexpanded for syncing.
 
         #: Parameters that Tune varies across searches.
-        self.evaluated_params = evaluated_params or []
+        self.evaluated_params = evaluated_params or {}
         self.experiment_tag = experiment_tag
         trainable_cls = self._get_trainable_cls()
         if trainable_cls and hasattr(trainable_cls,
@@ -237,6 +237,7 @@ class Trial(object):
             self.result_logger = UnifiedLogger(
                 self.config,
                 self.logdir,
+                trial=self,
                 loggers=self.loggers,
                 sync_function=self.sync_to_driver_fn)