Dataset learning

.gitignore

@@ -16,3 +16,6 @@ scripts/env.sh
 Dataset*
 .vscode
 .idea
+dataset
+cpps
+dataset_*

config.yaml.template

@@ -10,6 +10,15 @@ environment:
     dataset_path: "./Dataset_multi/"
     programs_file: "./multicomp.json"
     clean_files: True
+    json_dataset: 
+        # Path to the dataset file
+        path: "./Dataset_pickle/full_legality_annotations_solver.pkl"
+        # Path to the dataset cpp files
+        cpps_path: "./Dataset_multi"
+        # Path where to save the updated dataset (without the extension, it will be inferred by the dataset format)
+        path_to_save_dataset: "./Dataset_pickle/full_legality_annotations_solver_updated"
+        # Supported formats are available int he dataset_utilities module
+        dataset_format: "PICKLE"
 
 
 tiramisu:

evaluate.py

@@ -1,9 +1,11 @@
 import ray.rllib.agents.ppo as ppo
 import os
 import json
+
 # import hydra
 import argparse
 import ray
+
 # from hydra.core.config_store import ConfigStore
 from ray import tune
 from ray.rllib.models.catalog import ModelCatalog
@@ -13,9 +15,12 @@
 from rl_interface.model import TiramisuModelMult
 from utils.environment_variables import configure_env_variables
 from utils.global_ray_variables import Actor, GlobalVarActor
-from utils.rl_autoscheduler_config import (RLAutoSchedulerConfig,
-                                           dict_to_config, parse_yaml_file,
-                                           read_yaml_file)
+from utils.rl_autoscheduler_config import (
+    RLAutoSchedulerConfig,
+    dict_to_config,
+    parse_yaml_file,
+    read_yaml_file,
+)
 
 
 def get_arguments():
@@ -27,23 +32,23 @@ def get_arguments():
 
 # @hydra.main(config_path="config", config_name="config")
 def main(config: RLAutoSchedulerConfig, checkpoint=None):
-    if checkpoint is None: return
+    if checkpoint is None:
+        return
     configure_env_variables(config)
     best_checkpoint = os.path.join(config.ray.base_path, checkpoint)
     with ray.init(num_cpus=config.ray.ray_num_cpus):
         progs_list_registery = GlobalVarActor.remote(
             config.environment.programs_file,
             config.environment.dataset_path,
-            num_workers=config.ray.num_workers)
+            num_workers=config.ray.num_workers,
+        )
         shared_variable_actor = Actor.remote(progs_list_registery)
 
         register_env(
             "Tiramisu_env_v1",
-            lambda a: TiramisuScheduleEnvironment(config, shared_variable_actor
-                                                  ),
+            lambda a: TiramisuScheduleEnvironment(config, shared_variable_actor),
         )
-        ModelCatalog.register_custom_model("tiramisu_model_v1",
-                                           TiramisuModelMult)
+        ModelCatalog.register_custom_model("tiramisu_model_v1", TiramisuModelMult)
 
         agent = ppo.PPOTrainer(
             env="Tiramisu_env_v1",
@@ -62,7 +67,7 @@ def main(config: RLAutoSchedulerConfig, checkpoint=None):
                     "custom_model_config": {
                         "layer_sizes": list(config.model.layer_sizes),
                         "drops": list(config.model.drops),
-                    }
+                    },
                 },
             },
         )
@@ -87,7 +92,7 @@ def main(config: RLAutoSchedulerConfig, checkpoint=None):
                 except:
                     print("error", action, observation, reward, done)
                     continue
-            result["schedule_str"] = env.schedule_object.schedule_str
+            result["sched_str"] = env.schedule_object.sched_str
             result["speedup"] = env.schedule_controller.speedup
             results.append(result)
             with open("results.json", "w+") as file:

requirements.txt

@@ -1,7 +1,10 @@
 PyYAML == 6.0
-gym == 0.21.0
+gym==0.21.0
+gymnasium==0.27.1
 numpy == 1.23.1
-ray[rllib] == 1.13.0
+ray[rllib] == 2.2.0
 sympy == 1.10.1
 torch == 1.12.0
 tqdm == 4.64.0
+pandas == 1.5.3
+tensorflow_probability == 0.19.0

rl_interface/__init__.py

@@ -1,5 +0,0 @@
-from .action import *
-from .environment import *
-from .model import *
-from .reward import *
-from .utils import *