torch_geometric.data.lightning.LightningLinkData

class LightningLinkData(data: Union[Data, HeteroData], input_train_edges: Union[Tensor, None, Tuple[str, str, str], Tuple[Tuple[str, str, str], Optional[Tensor]]] = None, input_train_labels: Optional[Tensor] = None, input_train_time: Optional[Tensor] = None, input_val_edges: Union[Tensor, None, Tuple[str, str, str], Tuple[Tuple[str, str, str], Optional[Tensor]]] = None, input_val_labels: Optional[Tensor] = None, input_val_time: Optional[Tensor] = None, input_test_edges: Union[Tensor, None, Tuple[str, str, str], Tuple[Tuple[str, str, str], Optional[Tensor]]] = None, input_test_labels: Optional[Tensor] = None, input_test_time: Optional[Tensor] = None, input_pred_edges: Union[Tensor, None, Tuple[str, str, str], Tuple[Tuple[str, str, str], Optional[Tensor]]] = None, input_pred_labels: Optional[Tensor] = None, input_pred_time: Optional[Tensor] = None, loader: str = 'neighbor', link_sampler: Optional[BaseSampler] = None, eval_loader_kwargs: Optional[Dict[str, Any]] = None, **kwargs: Any)[source]

Bases: LightningData

Converts a Data or HeteroData object into a pytorch_lightning.LightningDataModule variant. It can then be automatically used as a datamodule for multi-GPU link-level training via PyTorch Lightning. LightningDataset will take care of providing mini-batches via LinkNeighborLoader.

Note

Currently only the pytorch_lightning.strategies.SingleDeviceStrategy and pytorch_lightning.strategies.DDPStrategy training strategies of PyTorch Lightning are supported in order to correctly share data across all devices/processes:

import pytorch_lightning as pl
trainer = pl.Trainer(strategy="ddp_spawn", accelerator="gpu",
                     devices=4)
trainer.fit(model, datamodule)
Parameters:

  • data (Data or HeteroData or Tuple[FeatureStore, GraphStore]) – The Data or HeteroData graph object, or a tuple of a FeatureStore and GraphStore objects.

  • input_train_edges (Tensor or EdgeType or Tuple[EdgeType, Tensor]) – The training edges. (default: None)

  • input_train_labels (torch.Tensor, optional) – The labels of training edges. (default: None)

  • input_train_time (torch.Tensor, optional) – The timestamp of training edges. (default: None)

  • input_val_edges (Tensor or EdgeType or Tuple[EdgeType, Tensor]) – The validation edges. (default: None)

  • input_val_labels (torch.Tensor, optional) – The labels of validation edges. (default: None)

  • input_val_time (torch.Tensor, optional) – The timestamp of validation edges. (default: None)

  • input_test_edges (Tensor or EdgeType or Tuple[EdgeType, Tensor]) – The test edges. (default: None)

  • input_test_labels (torch.Tensor, optional) – The labels of test edges. (default: None)

  • input_test_time (torch.Tensor, optional) – The timestamp of test edges. (default: None)

  • input_pred_edges (Tensor or EdgeType or Tuple[EdgeType, Tensor]) – The prediction edges. (default: None)

  • input_pred_labels (torch.Tensor, optional) – The labels of prediction edges. (default: None)

  • input_pred_time (torch.Tensor, optional) – The timestamp of prediction edges. (default: None)

  • loader (str) – The scalability technique to use ("full", "neighbor"). (default: "neighbor")

  • link_sampler (BaseSampler, optional) – A custom sampler object to generate mini-batches. If set, will ignore the loader option. (default: None)

  • eval_loader_kwargs (Dict[str, Any], optional) – Custom keyword arguments that override the torch_geometric.loader.LinkNeighborLoader configuration during evaluation. (default: None)

  • **kwargs (optional) – Additional arguments of torch_geometric.loader.LinkNeighborLoader.