openhgnn.dataset.AsLinkPredictionDataset

class AsLinkPredictionDataset(dataset, target_link, target_link_r, split_ratio=None, neg_ratio=3, neg_sampler='global', **kwargs)[源代码]

Repurpose a dataset for link prediction task.

The created dataset will include data needed for link prediction. It will keep only the first graph in the provided dataset and generate train/val/test edges according to the given split ratio, and the correspondent negative edges based on the neg_ratio. The generated edges will be cached to disk for fast re-loading. If the provided split ratio differs from the cached one, it will re-process the dataset properly.

参数:
  • dataset (DGLDataset) – The dataset to be converted.

  • split_ratio ((float, float, float), optional) – Split ratios for training, validation and test sets. Must sum to one.

  • neg_ratio (int, optional) – Indicate how much negative samples to be sampled The number of the negative samples will be equal or less than neg_ratio * num_positive_edges.

  • target_link (list[tuple[str, str, str]]) – The edge types on which predictions are make.

  • target_link_r (list[tuple[str, str, str]], optional) – The reverse edge types of the target links. Used to remove reverse edges of val/test edges from train graph.

  • neg_sampler (str, optional) – Indicate how negative edges of val/test edges are sampled. ‘global’ or ‘per_source’.

train_graph

The DGLHeteroGraph for training

Type:

DGLHeteroGraph

pos_val_graph

The DGLHeteroGraph containing positive validation edges

Type:

DGLHeteroGraph

pos_test_graph

The DGLHeteroGraph containing positive test edges

Type:

DGLHeteroGraph

neg_val_graph

The DGLHeteroGraph containing negative validation edges

Type:

DGLHeteroGraph

neg_test_graph

The DGLHeteroGraph containing negative test edges

Type:

DGLHeteroGraph