Update ax_multiobjective_nas_tutorial.py to use TensorboardMetric

.ci/docker/requirements.txt

@@ -31,7 +31,7 @@ torchx
 # TODO: use stable 0.5 when released
 -e git+https://github.com/pytorch/rl.git#egg=torchrl
 -e git+https://github.com/pytorch/tensordict.git#egg=tensordict
-ax-platform
+ax-platform>==0.4.0
 nbformat>==5.9.2
 datasets
 transformers
@@ -69,4 +69,4 @@ pygame==2.1.2
 pycocotools
 semilearn==0.3.2
 torchao==0.0.3
-segment_anything==1.0
+segment_anything==1.0

intermediate_source/ax_multiobjective_nas_tutorial.py

@@ -232,21 +232,21 @@ def trainer(
 # we get the logic to read and parse the TensorBoard logs for free.
 #
 
-from ax.metrics.tensorboard import TensorboardCurveMetric
+from ax.metrics.tensorboard import TensorboardMetric
+from tensorboard.backend.event_processing import plugin_event_multiplexer as event_multiplexer
 
-
-class MyTensorboardMetric(TensorboardCurveMetric):
+class MyTensorboardMetric(TensorboardMetric):
 
     # NOTE: We need to tell the new TensorBoard metric how to get the id /
     # file handle for the TensorBoard logs from a trial. In this case
     # our convention is to just save a separate file per trial in
     # the prespecified log dir.
-    @classmethod
-    def get_ids_from_trials(cls, trials):
-        return {
-            trial.index: Path(log_dir).joinpath(str(trial.index)).as_posix()
-            for trial in trials
-        }
+    def _get_event_multiplexer_for_trial(self, trial):
+        mul = event_multiplexer.EventMultiplexer(max_reload_threads=20)
+        mul.AddRunsFromDirectory(Path(log_dir).joinpath(str(trial.index)).as_posix(), None)
+        mul.Reload()
+
+        return mul
 
     # This indicates whether the metric is queryable while the trial is
     # still running. We don't use this in the current tutorial, but Ax
@@ -266,12 +266,12 @@ def is_available_while_running(cls):
 
 val_acc = MyTensorboardMetric(
     name="val_acc",
-    curve_name="val_acc",
+    tag="val_acc",
     lower_is_better=False,
 )
 model_num_params = MyTensorboardMetric(
     name="num_params",
-    curve_name="num_params",
+    tag="num_params",
     lower_is_better=True,
 )