add the namescope for monitor

google · Jul 26, 2021 · 9bada5d · 9bada5d
1 parent a4860ca
commit 9bada5d
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 9 deletions.
diff --git a/compiler_opt/rl/local_data_collector.py b/compiler_opt/rl/local_data_collector.py
@@ -98,7 +98,7 @@ def _schedule_jobs(self, policy_path, sampled_file_paths):
 
   def collect_data(
       self, policy_path: str
-  ) -> Tuple[Iterator[trajectory.Trajectory], Dict[str, float]]:
+  ) -> Tuple[Iterator[trajectory.Trajectory], Dict[str, Dict[str, float]]]:
     """Collect data for a given policy.
 
     Args:
@@ -166,7 +166,8 @@ def wait_for_termination():
     self._default_policy_size_map.update(
         {file_paths: res.get()[1] for (file_paths, res) in successful_work})
 
-    monitor_dict = {'success_modules': len(finished_work)}
+    monitor_dict = {}
+    monitor_dict['default'] = {'success_modules': len(finished_work)}
 
     return self._parser(sequence_examples), monitor_dict
 

diff --git a/compiler_opt/rl/local_data_collector_test.py b/compiler_opt/rl/local_data_collector_test.py
@@ -61,13 +61,13 @@ def _test_iterator_fn(data_list):
     data_iterator, monitor_dict = collector.collect_data(policy_path='policy')
     data = list(data_iterator)
     self.assertEqual([1, 2, 3], data)
-    expected_monitor_dict = {'success_modules': 10}
+    expected_monitor_dict = {'default': {'success_modules': 10}}
     self.assertEqual(expected_monitor_dict, monitor_dict)
 
     data_iterator, monitor_dict = collector.collect_data(policy_path='policy')
     data = list(data_iterator)
     self.assertEqual([4, 5, 6], data)
-    expected_monitor_dict = {'success_modules': 10}
+    expected_monitor_dict = {'default': {'success_modules': 10}}
     self.assertEqual(expected_monitor_dict, monitor_dict)
 
     collector.close_pool()

diff --git a/compiler_opt/rl/trainer.py b/compiler_opt/rl/trainer.py
@@ -120,7 +120,7 @@ def _update_metrics(self, experience, monitor_dict):
         experience.reward, sample_weight=is_action)
     self._num_trajectories.update_state(experience.is_first())
 
-    with tf.name_scope('Monitor/'):
+    with tf.name_scope('default/'):
       tf.summary.scalar(
           name='data_action_mean',
           data=self._data_action_mean.result(),
@@ -133,8 +133,11 @@ def _update_metrics(self, experience, monitor_dict):
           name='num_trajectories',
           data=self._num_trajectories.result(),
           step=self._global_step)
-      for key, value in monitor_dict.items():
-        tf.summary.scalar(name=key, data=value, step=self._global_step)
+
+    for name_scope, d in monitor_dict.items():
+      with tf.name_scope(name_scope + '/'):
+        for key, value in d.items():
+          tf.summary.scalar(name=key, data=value, step=self._global_step)
 
     tf.summary.histogram(
         name='reward', data=experience.reward, step=self._global_step)

diff --git a/compiler_opt/rl/trainer_test.py b/compiler_opt/rl/trainer_test.py
@@ -94,7 +94,7 @@ def test_training(self):
     self.assertEqual(0, test_trainer._global_step.numpy())
 
     dataset_iter = _create_test_data(batch_size=3, sequence_length=3)
-    monitor_dict = {'test': 1}
+    monitor_dict = {'default': {'test': 1}}
 
     with mock.patch.object(
         tf.summary, 'scalar', autospec=True) as mock_scalar_summary:
@@ -117,7 +117,7 @@ def test_training_with_multiple_times(self):
     self.assertEqual(0, test_trainer._global_step.numpy())
 
     dataset_iter = _create_test_data(batch_size=3, sequence_length=3)
-    monitor_dict = {'test': 1}
+    monitor_dict = {'default': {'test': 1}}
     test_trainer.train(dataset_iter, monitor_dict, num_iterations=10)
     self.assertEqual(10, test_trainer._global_step.numpy())