From f13b821c2a0bf56ce515986f8ca0416ef8df9a86 Mon Sep 17 00:00:00 2001 From: Oscar Barrera Date: Mon, 8 Jul 2024 14:57:05 -0400 Subject: [PATCH 1/4] Updated Examples to GraphNetDataModule --- examples/04_training/01_train_dynedge.py | 42 ++++++++++----- examples/04_training/02_train_tito_model.py | 53 ++++++++++++------- examples/04_training/05_train_RNN_TITO.py | 53 ++++++++++++------- examples/04_training/06_train_icemix_model.py | 53 ++++++++++++------- examples/04_training/README.md | 34 ++++++------ 5 files changed, 147 insertions(+), 88 deletions(-) diff --git a/examples/04_training/01_train_dynedge.py b/examples/04_training/01_train_dynedge.py index 3e1d08d2c..c61df789b 100644 --- a/examples/04_training/01_train_dynedge.py +++ b/examples/04_training/01_train_dynedge.py @@ -16,9 +16,11 @@ from graphnet.models.task.reconstruction import EnergyReconstruction from graphnet.training.callbacks import PiecewiseLinearLR from graphnet.training.loss_functions import LogCoshLoss -from graphnet.training.utils import make_train_validation_dataloader from graphnet.utilities.argparse import ArgumentParser from graphnet.utilities.logging import Logger +from graphnet.data import GraphNeTDataModule +from graphnet.data.dataset import SQLiteDataset +from graphnet.data.dataset import ParquetDataset # Constants features = FEATURES.PROMETHEUS @@ -68,6 +70,9 @@ def main( "gpus": gpus, "max_epochs": max_epochs, }, + "dataset_reference": SQLiteDataset + if path.endswith(".db") + else ParquetDataset, } archive = os.path.join(EXAMPLE_OUTPUT_DIR, "train_model_without_configs") @@ -79,21 +84,30 @@ def main( # Define graph representation graph_definition = KNNGraph(detector=Prometheus()) - ( - training_dataloader, - validation_dataloader, - ) = make_train_validation_dataloader( - db=config["path"], - graph_definition=graph_definition, - pulsemaps=config["pulsemap"], - features=features, - truth=truth, - batch_size=config["batch_size"], - num_workers=config["num_workers"], - truth_table=truth_table, - selection=None, + # Use GraphNetDataModule to load in data + dm = GraphNeTDataModule( + dataset_reference=config["dataset_reference"], + dataset_args={ + "truth": truth, + "truth_table": truth_table, + "features": features, + "graph_definition": graph_definition, + "pulsemaps": [config["pulsemap"]], + "path": config["path"], + }, + train_dataloader_kwargs={ + "batch_size": config["batch_size"], + "num_workers": config["num_workers"], + }, + test_dataloader_kwargs={ + "batch_size": config["batch_size"], + "num_workers": config["num_workers"], + }, ) + training_dataloader = dm.train_dataloader + validation_dataloader = dm.val_dataloader + # Building model backbone = DynEdge( diff --git a/examples/04_training/02_train_tito_model.py b/examples/04_training/02_train_tito_model.py index 2c6a1b4e0..50c5f7fd6 100644 --- a/examples/04_training/02_train_tito_model.py +++ b/examples/04_training/02_train_tito_model.py @@ -18,9 +18,11 @@ ) from graphnet.training.labels import Direction from graphnet.training.loss_functions import VonMisesFisher3DLoss -from graphnet.training.utils import make_train_validation_dataloader from graphnet.utilities.argparse import ArgumentParser from graphnet.utilities.logging import Logger +from graphnet.data import GraphNeTDataModule +from graphnet.data.dataset import SQLiteDataset +from graphnet.data.dataset import ParquetDataset # Constants features = FEATURES.PROMETHEUS @@ -70,6 +72,9 @@ def main( "gpus": gpus, "max_epochs": max_epochs, }, + "dataset_reference": SQLiteDataset + if path.endswith(".db") + else ParquetDataset, } graph_definition = KNNGraph(detector=Prometheus()) @@ -79,27 +84,37 @@ def main( # Log configuration to W&B wandb_logger.experiment.config.update(config) - ( - training_dataloader, - validation_dataloader, - ) = make_train_validation_dataloader( - db=config["path"], - graph_definition=graph_definition, - selection=None, - pulsemaps=config["pulsemap"], - features=features, - truth=truth, - batch_size=config["batch_size"], - num_workers=config["num_workers"], - truth_table=truth_table, - index_column="event_no", - labels={ - "direction": Direction( - azimuth_key="injection_azimuth", zenith_key="injection_zenith" - ) + # Use GraphNetDataModule to load in data + dm = GraphNeTDataModule( + dataset_reference=config["dataset_reference"], + dataset_args={ + "truth": truth, + "truth_table": truth_table, + "features": features, + "graph_definition": graph_definition, + "pulsemaps": [config["pulsemap"]], + "path": config["path"], + "index_column": "event_no", + "labels": { + "direction": Direction( + azimuth_key="injection_azimuth", + zenith_key="injection_zenith", + ) + }, + }, + train_dataloader_kwargs={ + "batch_size": config["batch_size"], + "num_workers": config["num_workers"], + }, + test_dataloader_kwargs={ + "batch_size": config["batch_size"], + "num_workers": config["num_workers"], }, ) + training_dataloader = dm.train_dataloader + validation_dataloader = dm.val_dataloader + # Building model backbone = DynEdgeTITO( nb_inputs=graph_definition.nb_outputs, diff --git a/examples/04_training/05_train_RNN_TITO.py b/examples/04_training/05_train_RNN_TITO.py index 26979fd5d..8e521bd1d 100644 --- a/examples/04_training/05_train_RNN_TITO.py +++ b/examples/04_training/05_train_RNN_TITO.py @@ -22,9 +22,11 @@ ) from graphnet.training.labels import Direction from graphnet.training.loss_functions import VonMisesFisher3DLoss -from graphnet.training.utils import make_train_validation_dataloader from graphnet.utilities.argparse import ArgumentParser from graphnet.utilities.logging import Logger +from graphnet.data import GraphNeTDataModule +from graphnet.data.dataset import SQLiteDataset +from graphnet.data.dataset import ParquetDataset # Constants features = FEATURES.PROMETHEUS @@ -74,6 +76,9 @@ def main( "gpus": gpus, "max_epochs": max_epochs, }, + "dataset_reference": SQLiteDataset + if path.endswith(".db") + else ParquetDataset, } graph_definition = KNNGraph( @@ -91,27 +96,37 @@ def main( # Log configuration to W&B wandb_logger.experiment.config.update(config) - ( - training_dataloader, - validation_dataloader, - ) = make_train_validation_dataloader( - db=config["path"], - graph_definition=graph_definition, - selection=None, - pulsemaps=config["pulsemap"], - features=features, - truth=truth, - batch_size=config["batch_size"], - num_workers=config["num_workers"], - truth_table=truth_table, - index_column="event_no", - labels={ - "direction": Direction( - azimuth_key="injection_azimuth", zenith_key="injection_zenith" - ) + # Use GraphNetDataModule to load in data + dm = GraphNeTDataModule( + dataset_reference=config["dataset_reference"], + dataset_args={ + "truth": truth, + "truth_table": truth_table, + "features": features, + "graph_definition": graph_definition, + "pulsemaps": [config["pulsemap"]], + "path": config["path"], + "index_column": "event_no", + "labels": { + "direction": Direction( + azimuth_key="injection_azimuth", + zenith_key="injection_zenith", + ) + }, + }, + train_dataloader_kwargs={ + "batch_size": config["batch_size"], + "num_workers": config["num_workers"], + }, + test_dataloader_kwargs={ + "batch_size": config["batch_size"], + "num_workers": config["num_workers"], }, ) + training_dataloader = dm.train_dataloader + validation_dataloader = dm.val_dataloader + # Building model backbone = RNN_TITO( nb_inputs=graph_definition.nb_outputs, diff --git a/examples/04_training/06_train_icemix_model.py b/examples/04_training/06_train_icemix_model.py index 450c4e379..6f8b58228 100644 --- a/examples/04_training/06_train_icemix_model.py +++ b/examples/04_training/06_train_icemix_model.py @@ -23,9 +23,11 @@ ) from graphnet.training.labels import Direction from graphnet.training.loss_functions import VonMisesFisher3DLoss -from graphnet.training.utils import make_train_validation_dataloader from graphnet.utilities.argparse import ArgumentParser from graphnet.utilities.logging import Logger +from graphnet.data import GraphNeTDataModule +from graphnet.data.dataset import SQLiteDataset +from graphnet.data.dataset import ParquetDataset # Constants features = FEATURES.PROMETHEUS @@ -76,6 +78,9 @@ def main( "max_epochs": max_epochs, "distribution_strategy": "ddp_find_unused_parameters_true", }, + "dataset_reference": SQLiteDataset + if path.endswith(".db") + else ParquetDataset, } graph_definition = KNNGraph( @@ -96,27 +101,37 @@ def main( # Log configuration to W&B wandb_logger.experiment.config.update(config) - ( - training_dataloader, - validation_dataloader, - ) = make_train_validation_dataloader( - db=config["path"], - graph_definition=graph_definition, - selection=None, - pulsemaps=config["pulsemap"], - features=features, - truth=truth, - batch_size=config["batch_size"], - num_workers=config["num_workers"], - truth_table=truth_table, - index_column="event_no", - labels={ - "direction": Direction( - azimuth_key="injection_azimuth", zenith_key="injection_zenith" - ) + # Use GraphNetDataModule to load in data + dm = GraphNeTDataModule( + dataset_reference=config["dataset_reference"], + dataset_args={ + "truth": truth, + "truth_table": truth_table, + "features": features, + "graph_definition": graph_definition, + "pulsemaps": [config["pulsemap"]], + "path": config["path"], + "index_column": "event_no", + "labels": { + "direction": Direction( + azimuth_key="injection_azimuth", + zenith_key="injection_zenith", + ) + }, + }, + train_dataloader_kwargs={ + "batch_size": config["batch_size"], + "num_workers": config["num_workers"], + }, + test_dataloader_kwargs={ + "batch_size": config["batch_size"], + "num_workers": config["num_workers"], }, ) + training_dataloader = dm.train_dataloader + validation_dataloader = dm.val_dataloader + # Building model backbone = DeepIce( hidden_dim=768, diff --git a/examples/04_training/README.md b/examples/04_training/README.md index 934767247..e94830cdf 100644 --- a/examples/04_training/README.md +++ b/examples/04_training/README.md @@ -2,7 +2,7 @@ This subfolder contains two main training scripts: -**`01_train_dynedge.py`** ** Shows how to train a GNN on neutrino telescope data **without configuration files,** i.e., by programatically constructing the dataset and model used. This is good for debugging and experimenting with different dataset settings and model configurations, as it is easier to build the model using the API than by writing configuration files from scratch. **This is our recommended way of getting started with the library**. For instance, try running: +**`01_train_dynedge.py`** Shows how to train a GNN on neutrino telescope data **without configuration files,** i.e., by programatically constructing the dataset and model used. This is good for debugging and experimenting with different dataset settings and model configurations, as it is easier to build the model using the API than by writing configuration files from scratch. **This is our recommended way of getting started with the library**. For instance, try running: ```bash # Show the CLI @@ -11,35 +11,35 @@ This subfolder contains two main training scripts: # Train energy regression model (graphnet) $ python examples/04_training/01_train_dynedge.py -# Same as above, as this is the default model config. -(graphnet) $ python examples/04_training/01_train_model.py \ - --model-config configs/models/example_energy_reconstruction_model.yml - # Train using a single GPU (graphnet) $ python examples/04_training/01_train_dynedge.py --gpus 0 # Train using multiple GPUs (graphnet) $ python examples/04_training/01_train_dynedge.py --gpus 0 1 +``` + +**`03_train_model_dynedge_from_config.py`** Shows how to train a GNN on neutrino telescope data **using configuration files** to construct the dataset that loads the data and the model that is trained. This is the recommended way to configure standard dataset and models, as it is easier to ready and share than doing so in pure code. This example can be run using a few different models targeting different physics use cases. For instance, you can try running: + +```bash +# Show the CLI +(graphnet) $ python examples/04_training/03_train_dynedge_from_config.py --help + +# Train energy regression model +(graphnet) $ python examples/04_training/03_train_dynedge_from_config.py +# Same as above, as this is the default model config. +(graphnet) $ python examples/04_training/03_train_dynedge_from_config.py \ + --model-config configs/models/example_energy_reconstruction_model.yml + # Train a vertex position reconstruction model -(graphnet) $ python examples/04_training/01_train_dynedge.py \ +(graphnet) $ python examples/04_training/03_train_dynedge_from_config.py \ --model-config configs/models/example_vertex_position_reconstruction_model.yml # Trains a direction (zenith, azimuth) reconstruction model. Note that the # chosen `Task` in the model config file also returns estimated "kappa" values, # i.e. inverse variance, for each predicted feature, meaning that we need to # manually specify the names of these. -(graphnet) $ python examples/04_training/01_train_model_dynedge.py --gpus 0 \ +(graphnet) $ python examples/04_training/03_train_dynedge_from_config.py --gpus 0 \ --model-config configs/models/example_direction_reconstruction_model.yml \ --prediction-names zenith_pred zenith_kappa_pred azimuth_pred azimuth_kappa_pred ``` - -**`03_train_model_dynedge_from_config.py** Shows how to train a GNN on neutrino telescope data **using configuration files** to construct the dataset that loads the data and the model that is trained. This is the recommended way to configure standard dataset and models, as it is easier to ready and share than doing so in pure code. This example can be run using a few different models targeting different physics use cases. For instance, you can try running: - -```bash -# Show the CLI -(graphnet) $ python examples/04_training/02_train_dynedge_from_config.py --help - -# Train energy regression model -(graphnet) $ python examples/04_training/02_train_dynedge_from_config.py -``` From b5331332a8904a217a6328b74cb937d4755edc9e Mon Sep 17 00:00:00 2001 From: Oscar Barrera Date: Mon, 8 Jul 2024 17:02:40 -0400 Subject: [PATCH 2/4] Added label argument to ParquetDataset --- src/graphnet/data/dataset/parquet/parquet_dataset.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/graphnet/data/dataset/parquet/parquet_dataset.py b/src/graphnet/data/dataset/parquet/parquet_dataset.py index 0186a7a72..be854bd28 100644 --- a/src/graphnet/data/dataset/parquet/parquet_dataset.py +++ b/src/graphnet/data/dataset/parquet/parquet_dataset.py @@ -54,6 +54,7 @@ def __init__( loss_weight_default_value: Optional[float] = None, seed: Optional[int] = None, cache_size: int = 1, + labels: Optional[Dict[str, Any]] = None, ): """Construct Dataset. @@ -122,6 +123,7 @@ def __init__( loss_weight_default_value=loss_weight_default_value, seed=seed, graph_definition=graph_definition, + labels = labels, ) # mypy.. From e78a8007ac6660fb8fc0439f7b286cb243a92c16 Mon Sep 17 00:00:00 2001 From: Oscar Barrera <168393473+OscarBarreraGithub@users.noreply.github.com> Date: Mon, 22 Jul 2024 16:56:14 -0400 Subject: [PATCH 3/4] Update parquet_dataset.py Removed spam before and after '=' in line 126 --- src/graphnet/data/dataset/parquet/parquet_dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/graphnet/data/dataset/parquet/parquet_dataset.py b/src/graphnet/data/dataset/parquet/parquet_dataset.py index be854bd28..dfc5720f1 100644 --- a/src/graphnet/data/dataset/parquet/parquet_dataset.py +++ b/src/graphnet/data/dataset/parquet/parquet_dataset.py @@ -123,7 +123,7 @@ def __init__( loss_weight_default_value=loss_weight_default_value, seed=seed, graph_definition=graph_definition, - labels = labels, + labels=labels, ) # mypy.. From 85a81ed2e19fc6932ecc7b21bd1d7b3953e41f76 Mon Sep 17 00:00:00 2001 From: Oscar Barrera <168393473+OscarBarreraGithub@users.noreply.github.com> Date: Mon, 22 Jul 2024 17:01:16 -0400 Subject: [PATCH 4/4] Update parquet_dataset.py Added 'labels' argument description --- src/graphnet/data/dataset/parquet/parquet_dataset.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/graphnet/data/dataset/parquet/parquet_dataset.py b/src/graphnet/data/dataset/parquet/parquet_dataset.py index dfc5720f1..84e938c44 100644 --- a/src/graphnet/data/dataset/parquet/parquet_dataset.py +++ b/src/graphnet/data/dataset/parquet/parquet_dataset.py @@ -103,6 +103,7 @@ def __init__( graph_definition: Method that defines the graph representation. cache_size: Number of batches to cache in memory. Must be at least 1. Defaults to 1. + labels: Dictionary of labels to be added to the dataset. """ self._validate_selection(selection) # Base class constructor