diff --git a/examples/qm9_hpo/qm9_deephyper.py b/examples/qm9_hpo/qm9_deephyper.py index 51031c27a..948968021 100644 --- a/examples/qm9_hpo/qm9_deephyper.py +++ b/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