From f8f6e0b04af625a81041e459ad5af1ea9959a7b9 Mon Sep 17 00:00:00 2001 From: Aiden Grossman Date: Wed, 18 Dec 2024 23:15:09 +0000 Subject: [PATCH] Use actual corpus in es_trainer_lib This patch makes es_trainer_lib load an actual corpus from the provided path rather than creating a dummy test corpus. Additionally, we also begin passing the replace_compilation_flags parameter, which is necessary for rewriting the paths to PGO profiles. --- compiler_opt/es/es_trainer_lib.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/compiler_opt/es/es_trainer_lib.py b/compiler_opt/es/es_trainer_lib.py index 21a220bf..15430f1a 100644 --- a/compiler_opt/es/es_trainer_lib.py +++ b/compiler_opt/es/es_trainer_lib.py @@ -70,6 +70,7 @@ @gin.configurable def train(additional_compilation_flags=(), delete_compilation_flags=(), + replace_compilation_flags=(), worker_class=None): """Train with ES.""" @@ -113,11 +114,11 @@ def train(additional_compilation_flags=(), logging.info("Parameter dimension: %s", initial_parameters.shape) logging.info("Initial parameters: %s", initial_parameters) - cps = corpus.create_corpus_for_testing( + cps = corpus.Corpus( location=_TRAIN_CORPORA.value, - elements=[corpus.ModuleSpec(name="smth", size=1)], additional_flags=additional_compilation_flags, - delete_flags=delete_compilation_flags) + delete_flags=delete_compilation_flags, + replace_flags=replace_compilation_flags) # Construct policy saver saved_policy = policy_utils.create_actor_policy(greedy=True)