Splitting Datasets¶
We provide multiple splitting methods for datasets.
Contents
ConsecutiveSplitter¶

class
dgllife.utils.
ConsecutiveSplitter
[source]¶ Split datasets with the input order.
The dataset is split without permutation, so the splitting is deterministic.

static
k_fold_split
(dataset, k=5, log=True)[source]¶ Split the dataset for kfold cross validation by taking consecutive chunks.
 Parameters
 Returns
Each element of the list represents a fold and is a 2tuple
(train_set, val_set)
.train_set
andval_set
also havelen(dataset)
anddataset[i]
behaviors. Return type
list of 2tuples

static
train_val_test_split
(dataset, frac_train=0.8, frac_val=0.1, frac_test=0.1)[source]¶ Split the dataset into three consecutive chunks for training, validation and test.
 Parameters
dataset – We assume
len(dataset)
gives the size for the dataset anddataset[i]
gives the ith datapoint.frac_train (float) – Fraction of data to use for training. By default, we set this to be 0.8, i.e. 80% of the dataset is used for training.
frac_val (float) – Fraction of data to use for validation. By default, we set this to be 0.1, i.e. 10% of the dataset is used for validation.
frac_test (float) – Fraction of data to use for test. By default, we set this to be 0.1, i.e. 10% of the dataset is used for test.
 Returns
Subsets for training, validation and test that also have
len(dataset)
anddataset[i]
behaviors Return type
list of length 3

static
RandomSplitter¶

class
dgllife.utils.
RandomSplitter
[source]¶ Randomly reorder datasets and then split them.
The dataset is split with permutation and the splitting is hence random.

static
k_fold_split
(dataset, k=5, random_state=None, log=True)[source]¶ Randomly permute the dataset and then split it for kfold cross validation by taking consecutive chunks.
 Parameters
dataset – We assume
len(dataset)
gives the size for the dataset anddataset[i]
gives the ith datapoint.k (int) – Number of folds to use and should be no smaller than 2. Default to be 5.
random_state (None, int or array_like, optional) – Random seed used to initialize the pseudorandom number generator. Can be any integer between 0 and 2**32  1 inclusive, an array (or other sequence) of such integers, or None (the default). If seed is None, then RandomState will try to read data from /dev/urandom (or the Windows analogue) if available or seed from the clock otherwise.
log (bool) – Whether to print a message at the start of preparing each fold. Default to True.
 Returns
Each element of the list represents a fold and is a 2tuple
(train_set, val_set)
.train_set
andval_set
also havelen(dataset)
anddataset[i]
behaviors. Return type
list of 2tuples

static
train_val_test_split
(dataset, frac_train=0.8, frac_val=0.1, frac_test=0.1, random_state=None)[source]¶ Randomly permute the dataset and then split it into three consecutive chunks for training, validation and test.
 Parameters
dataset – We assume
len(dataset)
gives the size for the dataset anddataset[i]
gives the ith datapoint.frac_train (float) – Fraction of data to use for training. By default, we set this to be 0.8, i.e. 80% of the dataset is used for training.
frac_val (float) – Fraction of data to use for validation. By default, we set this to be 0.1, i.e. 10% of the dataset is used for validation.
frac_test (float) – Fraction of data to use for test. By default, we set this to be 0.1, i.e. 10% of the dataset is used for test.
random_state (None, int or array_like, optional) – Random seed used to initialize the pseudorandom number generator. Can be any integer between 0 and 2**32  1 inclusive, an array (or other sequence) of such integers, or None (the default). If seed is None, then RandomState will try to read data from /dev/urandom (or the Windows analogue) if available or seed from the clock otherwise.
 Returns
Subsets for training, validation and test, which also have
len(dataset)
anddataset[i]
behaviors. Return type
list of length 3

static
MolecularWeightSplitter¶

class
dgllife.utils.
MolecularWeightSplitter
[source]¶ Sort molecules based on their weights and then split them.

static
k_fold_split
(dataset, mols=None, sanitize=True, k=5, log_every_n=1000)[source]¶ Sort molecules based on their weights and then split them for kfold cross validation by taking consecutive chunks.
 Parameters
dataset – We assume
len(dataset)
gives the size for the dataset,dataset[i]
gives the ith datapoint anddataset.smiles[i]
gives the SMILES for the ith datapoint.mols (None or list of rdkit.Chem.rdchem.Mol) – None or precomputed RDKit molecule instances. If not None, we expect a oneonone correspondence between
dataset.smiles
andmols
, i.e.mols[i]
corresponds todataset.smiles[i]
. Default to None.sanitize (bool) – This argument only comes into effect when
mols
is None and decides whether sanitization is performed in initializing RDKit molecule instances. See https://www.rdkit.org/docs/RDKit_Book.html for details of the sanitization. Default to be True.k (int) – Number of folds to use and should be no smaller than 2. Default to be 5.
log_every_n (None or int) – Molecule related computation can take a long time for a large dataset and we want to learn the progress of processing. This can be done by printing a message whenever a batch of
log_every_n
molecules have been processed. If None, no messages will be printed. Default to 1000.
 Returns
Each element of the list represents a fold and is a 2tuple
(train_set, val_set)
.train_set
andval_set
also havelen(dataset)
anddataset[i]
behaviors. Return type
list of 2tuples

static
train_val_test_split
(dataset, mols=None, sanitize=True, frac_train=0.8, frac_val=0.1, frac_test=0.1, log_every_n=1000)[source]¶ Sort molecules based on their weights and then split them into three consecutive chunks for training, validation and test.
 Parameters
dataset – We assume
len(dataset)
gives the size for the dataset,dataset[i]
gives the ith datapoint anddataset.smiles[i]
gives the SMILES for the ith datapoint.mols (None or list of rdkit.Chem.rdchem.Mol) – None or precomputed RDKit molecule instances. If not None, we expect a oneonone correspondence between
dataset.smiles
andmols
, i.e.mols[i]
corresponds todataset.smiles[i]
. Default to None.sanitize (bool) – This argument only comes into effect when
mols
is None and decides whether sanitization is performed in initializing RDKit molecule instances. See https://www.rdkit.org/docs/RDKit_Book.html for details of the sanitization. Default to be True.frac_train (float) – Fraction of data to use for training. By default, we set this to be 0.8, i.e. 80% of the dataset is used for training.
frac_val (float) – Fraction of data to use for validation. By default, we set this to be 0.1, i.e. 10% of the dataset is used for validation.
frac_test (float) – Fraction of data to use for test. By default, we set this to be 0.1, i.e. 10% of the dataset is used for test.
log_every_n (None or int) – Molecule related computation can take a long time for a large dataset and we want to learn the progress of processing. This can be done by printing a message whenever a batch of
log_every_n
molecules have been processed. If None, no messages will be printed. Default to 1000.
 Returns
Subsets for training, validation and test, which also have
len(dataset)
anddataset[i]
behaviors Return type
list of length 3

static
ScaffoldSplitter¶

class
dgllife.utils.
ScaffoldSplitter
[source]¶ Group molecules based on their BemisMurcko scaffolds and then split the groups.
Group molecules so that all molecules in a group have a same scaffold (see reference). The dataset is then split at the level of groups.
References
 Bemis, G. W.; Murcko, M. A. “The Properties of Known Drugs.
Molecular Frameworks.” J. Med. Chem. 39:288793 (1996).

static
k_fold_split
(dataset, mols=None, sanitize=True, k=5, log_every_n=1000, scaffold_func='decompose')[source]¶ Group molecules based on their scaffolds and sort groups based on their sizes. The groups are then split for kfold cross validation.
Same as usual kfold splitting methods, each molecule will appear only once in the validation set among all folds. In addition, this method ensures that molecules with a same scaffold will be collectively in either the training set or the validation set for each fold.
Note that the folds can be highly imbalanced depending on the scaffold distribution in the dataset.
 Parameters
dataset – We assume
len(dataset)
gives the size for the dataset,dataset[i]
gives the ith datapoint anddataset.smiles[i]
gives the SMILES for the ith datapoint.mols (None or list of rdkit.Chem.rdchem.Mol) – None or precomputed RDKit molecule instances. If not None, we expect a oneonone correspondence between
dataset.smiles
andmols
, i.e.mols[i]
corresponds todataset.smiles[i]
. Default to None.sanitize (bool) – This argument only comes into effect when
mols
is None and decides whether sanitization is performed in initializing RDKit molecule instances. See https://www.rdkit.org/docs/RDKit_Book.html for details of the sanitization. Default to True.k (int) – Number of folds to use and should be no smaller than 2. Default to be 5.
log_every_n (None or int) – Molecule related computation can take a long time for a large dataset and we want to learn the progress of processing. This can be done by printing a message whenever a batch of
log_every_n
molecules have been processed. If None, no messages will be printed. Default to 1000.scaffold_func (str) – The function to use for computing scaffolds, which can be ‘decompose’ for using rdkit.Chem.AllChem.MurckoDecompose or ‘smiles’ for using rdkit.Chem.Scaffolds.MurckoScaffold.MurckoScaffoldSmiles.
 Returns
Each element of the list represents a fold and is a 2tuple
(train_set, val_set)
.train_set
andval_set
also havelen(dataset)
anddataset[i]
behaviors. Return type
list of 2tuples

static
train_val_test_split
(dataset, mols=None, sanitize=True, frac_train=0.8, frac_val=0.1, frac_test=0.1, log_every_n=1000, scaffold_func='decompose')[source]¶ Split the dataset into training, validation and test set based on molecular scaffolds.
This spliting method ensures that molecules with a same scaffold will be collectively in only one of the training, validation or test set. As a result, the fraction of dataset to use for training and validation tend to be smaller than
frac_train
andfrac_val
, while the fraction of dataset to use for test tends to be larger thanfrac_test
. Parameters
dataset – We assume
len(dataset)
gives the size for the dataset,dataset[i]
gives the ith datapoint anddataset.smiles[i]
gives the SMILES for the ith datapoint.mols (None or list of rdkit.Chem.rdchem.Mol) – None or precomputed RDKit molecule instances. If not None, we expect a oneonone correspondence between
dataset.smiles
andmols
, i.e.mols[i]
corresponds todataset.smiles[i]
. Default to None.sanitize (bool) – This argument only comes into effect when
mols
is None and decides whether sanitization is performed in initializing RDKit molecule instances. See https://www.rdkit.org/docs/RDKit_Book.html for details of the sanitization. Default to True.frac_train (float) – Fraction of data to use for training. By default, we set this to be 0.8, i.e. 80% of the dataset is used for training.
frac_val (float) – Fraction of data to use for validation. By default, we set this to be 0.1, i.e. 10% of the dataset is used for validation.
frac_test (float) – Fraction of data to use for test. By default, we set this to be 0.1, i.e. 10% of the dataset is used for test.
log_every_n (None or int) – Molecule related computation can take a long time for a large dataset and we want to learn the progress of processing. This can be done by printing a message whenever a batch of
log_every_n
molecules have been processed. If None, no messages will be printed. Default to 1000.scaffold_func (str) – The function to use for computing scaffolds, which can be ‘decompose’ for using rdkit.Chem.AllChem.MurckoDecompose or ‘smiles’ for using rdkit.Chem.Scaffolds.MurckoScaffold.MurckoScaffoldSmiles.
 Returns
Subsets for training, validation and test, which also have
len(dataset)
anddataset[i]
behaviors Return type
list of length 3
SingleTaskStratifiedSplitter¶

class
dgllife.utils.
SingleTaskStratifiedSplitter
[source]¶ Splits the dataset by stratification on a single task.
We sort the molecules based on their label values for a task and then repeatedly take buckets of datapoints to augment the training, validation and test subsets.

static
k_fold_split
(dataset, labels, task_id, k=5, log=True)[source]¶ Sort molecules based on their label values for a task and then split them for kfold cross validation by taking consecutive chunks.
 Parameters
dataset – We assume
len(dataset)
gives the size for the dataset,dataset[i]
gives the ith datapoint anddataset.smiles[i]
gives the SMILES for the ith datapoint.labels (tensor of shape (N, T)) – Dataset labels all tasks. N for the number of datapoints and T for the number of tasks.
task_id (int) – Index for the task.
k (int) – Number of folds to use and should be no smaller than 2. Default to be 5.
log (bool) – Whether to print a message at the start of preparing each fold.
 Returns
Each element of the list represents a fold and is a 2tuple
(train_set, val_set)
.train_set
andval_set
also havelen(dataset)
anddataset[i]
behaviors. Return type
list of 2tuples

static
train_val_test_split
(dataset, labels, task_id, frac_train=0.8, frac_val=0.1, frac_test=0.1, bucket_size=10, random_state=None)[source]¶ Split the dataset into training, validation and test subsets as stated above.
 Parameters
dataset – We assume
len(dataset)
gives the size for the dataset,dataset[i]
gives the ith datapoint anddataset.smiles[i]
gives the SMILES for the ith datapoint.labels (tensor of shape (N, T)) – Dataset labels all tasks. N for the number of datapoints and T for the number of tasks.
task_id (int) – Index for the task.
frac_train (float) – Fraction of data to use for training. By default, we set this to be 0.8, i.e. 80% of the dataset is used for training.
frac_val (float) – Fraction of data to use for validation. By default, we set this to be 0.1, i.e. 10% of the dataset is used for validation.
frac_test (float) – Fraction of data to use for test. By default, we set this to be 0.1, i.e. 10% of the dataset is used for test.
bucket_size (int) – Size of bucket of datapoints. Default to 10.
random_state (None, int or array_like, optional) – Random seed used to initialize the pseudorandom number generator. Can be any integer between 0 and 2**32  1 inclusive, an array (or other sequence) of such integers, or None (the default). If seed is None, then RandomState will try to read data from /dev/urandom (or the Windows analogue) if available or seed from the clock otherwise.
 Returns
Subsets for training, validation and test, which also have
len(dataset)
anddataset[i]
behaviors Return type
list of length 3

static