torch_geometric.data.lightning.LightningNodeData
- class LightningNodeData(data: Union[Data, HeteroData], input_train_nodes: Union[Tensor, None, str, Tuple[str, Optional[Tensor]]] = None, input_train_time: Optional[Tensor] = None, input_val_nodes: Union[Tensor, None, str, Tuple[str, Optional[Tensor]]] = None, input_val_time: Optional[Tensor] = None, input_test_nodes: Union[Tensor, None, str, Tuple[str, Optional[Tensor]]] = None, input_test_time: Optional[Tensor] = None, input_pred_nodes: Union[Tensor, None, str, Tuple[str, Optional[Tensor]]] = None, input_pred_time: Optional[Tensor] = None, loader: str = 'neighbor', node_sampler: Optional[BaseSampler] = None, eval_loader_kwargs: Optional[Dict[str, Any]] = None, **kwargs)[source]
Bases:
LightningData
Converts a
Data
orHeteroData
object into apytorch_lightning.LightningDataModule
variant. It can then be automatically used as adatamodule
for multi-GPU node-level training via PyTorch Lightning.LightningDataset
will take care of providing mini-batches viaNeighborLoader
.Note
Currently only the
pytorch_lightning.strategies.SingleDeviceStrategy
andpytorch_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) – The
Data
orHeteroData
graph object.input_train_nodes (torch.Tensor or str or (str, torch.Tensor)) – The indices of training nodes. If not given, will try to automatically infer them from the
data
object by searching fortrain_mask
,train_idx
, ortrain_index
attributes. (default:None
)input_train_time (torch.Tensor, optional) – The timestamp of training nodes. (default:
None
)input_val_nodes (torch.Tensor or str or (str, torch.Tensor)) – The indices of validation nodes. If not given, will try to automatically infer them from the
data
object by searching forval_mask
,valid_mask
,val_idx
,valid_idx
,val_index
, orvalid_index
attributes. (default:None
)input_val_time (torch.Tensor, optional) – The timestamp of validation edges. (default:
None
)input_test_nodes (torch.Tensor or str or (str, torch.Tensor)) – The indices of test nodes. If not given, will try to automatically infer them from the
data
object by searching fortest_mask
,test_idx
, ortest_index
attributes. (default:None
)input_test_time (torch.Tensor, optional) – The timestamp of test nodes. (default:
None
)input_pred_nodes (torch.Tensor or str or (str, torch.Tensor)) – The indices of prediction nodes. If not given, will try to automatically infer them from the
data
object by searching forpred_mask
,pred_idx
, orpred_index
attributes. (default:None
)input_pred_time (torch.Tensor, optional) – The timestamp of prediction nodes. (default:
None
)loader (str) – The scalability technique to use (
"full"
,"neighbor"
). (default:"neighbor"
)node_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.NeighborLoader
configuration during evaluation. (default:None
)**kwargs (optional) – Additional arguments of
torch_geometric.loader.NeighborLoader
.