add_nodepred_splitΒΆ
-
class
dgl.data.utils.
add_nodepred_split
(dataset, ratio, ntype=None)[source]ΒΆ Split the given dataset into training, validation and test sets for transductive node predction task.
It adds three node mask arrays
'train_mask'
,'val_mask'
and'test_mask'
, to each graph in the dataset. Each sample in the dataset thus must be aDGLGraph
.Fix the random seed of NumPy to make the result deterministic:
numpy.random.seed(42)
- Parameters
dataset (DGLDataset) β The dataset to modify.
ratio ((float, float, float)) β Split ratios for training, validation and test sets. Must sum to one.
ntype (str, optional) β The node type to add mask for.
Examples
>>> dataset = dgl.data.AmazonCoBuyComputerDataset() >>> print('train_mask' in dataset[0].ndata) False >>> dgl.data.utils.add_nodepred_split(dataset, [0.8, 0.1, 0.1]) >>> print('train_mask' in dataset[0].ndata) True