Minor changes

compiler_opt/es/blackbox_learner.py

@@ -196,6 +196,10 @@ def _log_tf_summary(self, rewards: List[float]) -> None:
  with self._summary_writer.as_default():
  tf.summary.scalar(
  'reward/average_reward_train', np.mean(rewards), step=self._step)
+
+ tf.summary.scalar(
+ 'reward/maximum_reward_train', np.max(rewards), step=self._step
+ )
 
  tf.summary.histogram('reward/reward_train', rewards, step=self._step)
 

compiler_opt/es/es_trainer_lib.py

@@ -21,6 +21,7 @@
 import gin
 import tensorflow as tf
 import os
+import shutil
 
 from compiler_opt.distributed import worker
 from compiler_opt.distributed.local import local_worker_manager
@@ -79,6 +80,7 @@ def __init__(self, *, all_gin):
  self._clang_path = '/usr/local/google/home/aidengrossman/programming/test_traces/clang'
  self._trace_path = '/usr/local/google/home/aidengrossman/programming/test_traces/execution_trace.pb'
  self._bb_trace_model_path = '/usr/local/google/home/aidengrossman/programming/test_traces/basic_block_trace_model'
+ self._models_for_test_path = '/usr/local/google/home/aidengrossman/programming/output_traces/'
 
  def es_compile(self, params: list[float], baseline_score: float) -> float:
  with tempfile.TemporaryDirectory() as tempdir:
@@ -99,8 +101,13 @@ def es_compile(self, params: list[float], baseline_score: float) -> float:
  self._clang_path, tflitedir)
  score = trace_data_collector.evaluate_compiled_corpus(
  tempdir, self._trace_path, self._bb_trace_model_path)
- print(score)
+
+ reward = compilation_runner._calculate_reward(score, baseline_score)
+ print(reward)
 
+ output_path = os.path.join(self._models_for_test_path, "model" + str(reward))
+ if reward > 0.2 and not os.path.exists(output_path):
+ shutil.copytree(tflitedir, output_path)
  return compilation_runner._calculate_reward(score, baseline_score)
 
 

compiler_opt/es/gin_configs/blackbox_learner.gin

@@ -3,8 +3,8 @@ import compiler_opt.rl.gin_external_configurables
 import compiler_opt.es.blackbox_optimizers
 
 # Blackbox learner config
-BlackboxLearnerConfig.total_steps = 1
-BlackboxLearnerConfig.total_num_perturbations = 5
+BlackboxLearnerConfig.total_steps = 100
+BlackboxLearnerConfig.total_num_perturbations = 100
 BlackboxLearnerConfig.blackbox_optimizer = %blackbox_optimizers.Algorithm.MONTE_CARLO
 BlackboxLearnerConfig.est_type = %blackbox_optimizers.EstimatorType.ANTITHETIC
 # BlackboxLearnerConfig.est_type = %blackbox_optimizers.EstimatorType.FORWARD_FD