graphnet.data.datamodule module

Base Dataloader class(es) used in graphnet.

class graphnet.data.datamodule.GraphNeTDataModule(dataset_reference, dataset_args, selection, test_selection, train_dataloader_kwargs, validation_dataloader_kwargs, test_dataloader_kwargs, train_val_split=[0.9, 0.1], split_seed)[source]

Bases: LightningDataModule, Logger

General Class for DataLoader Construction.

Create dataloaders from dataset.

Parameters:

  • dataset_reference (Union[Type[SQLiteDataset], Type[ParquetDataset], Type[Dataset]]) – A non-instantiated reference to the dataset class.

  • dataset_args (Dict[str, Any]) – Arguments to instantiate graphnet.data.dataset.Dataset with.

  • selection (Union[List[int], List[List[int]], None], default: None) – (Optional) a list of event id’s used for training and validation, Default None.

  • test_selection (Union[List[int], List[List[int]], None], default: None) – (Optional) a list of event id’s used for testing, Defaults to None.

  • train_dataloader_kwargs (Optional[Dict[str, Any]], default: None) – Arguments for the training DataLoader, Defaults{“batch_size”: 2, “num_workers”: 1}.

  • validation_dataloader_kwargs (Optional[Dict[str, Any]], default: None) – Arguments for the validation DataLoader. Defaults to train_dataloader_kwargs.

  • test_dataloader_kwargs (Optional[Dict[str, Any]], default: None) – Arguments for the test DataLoader, Defaults to train_dataloader_kwargs.

  • train_val_split (Optional) – Split ratio for training and validation sets. Default is [0.9, 0.10].

  • split_seed (int, default: 42) – seed used for shuffling and splitting selections into train/validation, Default 42.

prepare_data()[source]

Prepare the dataset for training.

Return type:

None

setup(stage)[source]

Prepare Datasets for DataLoaders.

Parameters:

stage (str) – lightning stage. Either “fit, validate, test, predict”

Return type:

None

property train_dataloader: DataLoader

Prepare and return the training DataLoader.

Returns:

The DataLoader configured for training.

Return type:

DataLoader

property val_dataloader: DataLoader

Prepare and return the validation DataLoader.

Returns:

The DataLoader configured for validation.

Return type:

DataLoader

property test_dataloader: DataLoader

Prepare and return the test DataLoader.

Returns:

The DataLoader configured for testing.

Return type:

DataLoader

teardown()[source]

Perform any necessary cleanup or shutdown procedures.

This method can be used for tasks such as closing SQLite connections after training. Override this method as needed.

Return type:

None

Returns:

None