Refactored train_eval to use num_modules from gin configs (#53)

Author: boomanaiden154

compiler_opt/rl/inlining/gin_configs/ppo_nn_agent.gin

@@ -12,6 +12,7 @@ include 'compiler_opt/rl/inlining/gin_configs/common.gin'
 train_eval.agent_name=%constant.AgentName.PPO
 train_eval.warmstart_policy_dir=''
 train_eval.num_policy_iterations=3000
+train_eval.num_modules=100
 train_eval.num_iterations=300
 train_eval.batch_size=256
 train_eval.train_sequence_length=16

compiler_opt/rl/regalloc/gin_configs/ppo_nn_agent.gin

@@ -13,6 +13,7 @@ include 'compiler_opt/rl/regalloc/gin_configs/network.gin'
 train_eval.agent_name=%constant.AgentName.PPO
 train_eval.warmstart_policy_dir=''
 train_eval.num_policy_iterations=3000
+train_eval.num_modules=512
 train_eval.num_iterations=200
 train_eval.batch_size=256
 train_eval.train_sequence_length=16

compiler_opt/rl/train_locally.py

@@ -48,8 +48,6 @@
 flags.DEFINE_integer(
     'num_workers', None,
     'Number of parallel data collection workers. `None` for max available')
-flags.DEFINE_integer('num_modules', 100,
-                     'Number of modules to collect data for each iteration.')
 flags.DEFINE_multi_string('gin_files', [],
                           'List of paths to gin configuration files.')
 flags.DEFINE_multi_string(
@@ -63,6 +61,7 @@
 def train_eval(agent_name=constant.AgentName.PPO,
                warmstart_policy_dir=None,
                num_policy_iterations=0,
+               num_modules=100,
                num_iterations=100,
                batch_size=64,
                train_sequence_length=1,
@@ -149,7 +148,7 @@ def sequence_example_iterator_fn(seq_ex: List[str]):
       delete_flags=delete_compilation_flags) as worker_pool:
     data_collector = local_data_collector.LocalDataCollector(
         file_paths=tuple(file_paths),
-        num_modules=FLAGS.num_modules,
+        num_modules=num_modules,
         worker_pool=worker_pool,
         parser=sequence_example_iterator_fn,
         reward_stat_map=reward_stat_map)