deephyper example with qm9 fixed

examples/qm9_hpo/qm9_deephyper.py

@@ -26,22 +26,40 @@ def qm9_pre_transform(data):
     return data
 
 
+log_name = "qm9"
+
+# Use built-in torch_geometric datasets.
+# Filter function above used to run quick example.
+# NOTE: data is moved to the device in the pre-transform.
+# NOTE: transforms/filters will NOT be re-run unless the qm9/processed/ directory is removed.
+dataset = torch_geometric.datasets.QM9(
+    root="dataset/qm9", pre_transform=qm9_pre_transform
+)
+
+# Configurable run choices (JSON file that accompanies this example script).
+filename = os.path.join(os.path.dirname(os.path.abspath(__file__)), "qm9.json")
+with open(filename, "r") as f:
+    config = json.load(f)
+verbosity = config["Verbosity"]["level"]
+
+trainset, valset, testset = hydragnn.preprocess.split_dataset(dataset, 0.8, False)
+(train_loader, val_loader, test_loader) = hydragnn.preprocess.create_dataloaders(
+    trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
+)
+
+
 def run(trial):
 
-    global config, log_name, train_loader, val_loader, test_loader
+    global config
 
-    # Configurable run choices (JSON file that accompanies this example script).
-    filename = os.path.join(os.path.dirname(os.path.abspath(__file__)), "qm9.json")
-    with open(filename, "r") as f:
-        config = json.load(f)
-    verbosity = config["Verbosity"]["level"]
+    trial_config = config
 
     # Always initialize for multi-rank training.
-    comm_size, rank = hydragnn.utils.setup_ddp()
+    comm_size, rank = hydragnn.utils.distributed.setup_ddp()
 
-    log_name = log_name + "_" + str(trial.id)
-    hydragnn.utils.setup_log(log_name)
-    writer = hydragnn.utils.get_summary_writer(log_name)
+    trial_log_name = log_name + "_" + str(trial.id)
+    hydragnn.utils.print.print_utils.setup_log(trial_log_name)
+    writer = hydragnn.utils.model.model.get_summary_writer(trial_log_name)
 
     ## Set up logging
     logging.basicConfig(
@@ -53,13 +71,13 @@ def run(trial):
     # log("Command: {0}\n".format(" ".join([x for x in sys.argv])), rank=0)
 
     # Update the config dictionary with the suggested hyperparameters
-    config["NeuralNetwork"]["Architecture"]["model_type"] = trial.parameters[
+    trial_config["NeuralNetwork"]["Architecture"]["model_type"] = trial.parameters[
         "model_type"
     ]
-    config["NeuralNetwork"]["Architecture"]["hidden_dim"] = trial.parameters[
+    trial_config["NeuralNetwork"]["Architecture"]["hidden_dim"] = trial.parameters[
         "hidden_dim"
     ]
-    config["NeuralNetwork"]["Architecture"]["num_conv_layers"] = trial.parameters[
+    trial_config["NeuralNetwork"]["Architecture"]["num_conv_layers"] = trial.parameters[
         "num_conv_layers"
     ]
 
@@ -68,35 +86,39 @@ def run(trial):
         for i in range(trial.parameters["num_headlayers"])
     ]
 
-    for head_type in config["NeuralNetwork"]["Architecture"]["output_heads"]:
-        config["NeuralNetwork"]["Architecture"]["output_heads"][head_type][
+    for head_type in trial_config["NeuralNetwork"]["Architecture"]["output_heads"]:
+        trial_config["NeuralNetwork"]["Architecture"]["output_heads"][head_type][
             "num_headlayers"
         ] = trial.parameters["num_headlayers"]
-        config["NeuralNetwork"]["Architecture"]["output_heads"][head_type][
+        trial_config["NeuralNetwork"]["Architecture"]["output_heads"][head_type][
             "dim_headlayers"
         ] = dim_headlayers
 
     if trial.parameters["model_type"] not in ["EGNN", "SchNet", "DimeNet"]:
-        config["NeuralNetwork"]["Architecture"]["equivariance"] = False
+        trial_config["NeuralNetwork"]["Architecture"]["equivariance"] = False
 
-    config = hydragnn.utils.update_config(config, train_loader, val_loader, test_loader)
+    trial_config = hydragnn.utils.input_config_parsing.update_config(
+        trial_config, train_loader, val_loader, test_loader
+    )
 
-    hydragnn.utils.save_config(config, log_name)
+    hydragnn.utils.input_config_parsing.save_config(trial_config, trial_log_name)
 
     model = hydragnn.models.create_model_config(
-        config=config["NeuralNetwork"],
+        config=trial_config["NeuralNetwork"],
         verbosity=verbosity,
     )
-    model = hydragnn.utils.get_distributed_model(model, verbosity)
+    model = hydragnn.utils.distributed.get_distributed_model(model, verbosity)
 
-    learning_rate = config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"]
+    learning_rate = trial_config["NeuralNetwork"]["Training"]["Optimizer"][
+        "learning_rate"
+    ]
     optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
     scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
         optimizer, mode="min", factor=0.5, patience=5, min_lr=0.00001
     )
 
-    hydragnn.utils.load_existing_model_config(
-        model, config["NeuralNetwork"]["Training"], optimizer=optimizer
+    hydragnn.utils.model.model.load_existing_model_config(
+        model, trial_config["NeuralNetwork"]["Training"], optimizer=optimizer
     )
 
     ##################################################################################################################
@@ -108,14 +130,14 @@ def run(trial):
         test_loader,
         writer,
         scheduler,
-        config["NeuralNetwork"],
-        log_name,
+        trial_config["NeuralNetwork"],
+        trial_log_name,
         verbosity,
         create_plots=False,
     )
 
-    hydragnn.utils.save_model(model, optimizer, log_name)
-    hydragnn.utils.print_timers(verbosity)
+    hydragnn.utils.model.model.save_model(model, optimizer, trial_log_name)
+    hydragnn.utils.print.print_distributed(verbosity)
 
     # Return the metric to minimize (e.g., validation loss)
     validation_loss, tasks_loss = hydragnn.train.validate(
@@ -133,24 +155,9 @@ def run(trial):
 
 if __name__ == "__main__":
 
-    log_name = "qm9"
-
-    # Use built-in torch_geometric datasets.
-    # Filter function above used to run quick example.
-    # NOTE: data is moved to the device in the pre-transform.
-    # NOTE: transforms/filters will NOT be re-run unless the qm9/processed/ directory is removed.
-    dataset = torch_geometric.datasets.QM9(
-        root="dataset/qm9", pre_transform=qm9_pre_transform
-    )
-
-    trainset, valset, testset = hydragnn.preprocess.split_dataset(dataset, 0.8, False)
-    (train_loader, val_loader, test_loader) = hydragnn.preprocess.create_dataloaders(
-        trainset, valset, testset, 64
-    )
-
     # Choose the sampler (e.g., TPESampler or RandomSampler)
-    from deephyper.problem import HpProblem
-    from deephyper.search.hps import CBO
+    from deephyper.hpo import HpProblem, CBO
+    from deephyper.hpo import CBO
     from deephyper.evaluator import Evaluator
 
     # define the variable you want to optimize