gigl.distributed.dist_link_prediction_dataset#
Attributes#
Classes#
This class is inherited from GraphLearn-for-PyTorch's DistDataset class. We override the __init__ functionality to support positive and |
Module Contents#
- class gigl.distributed.dist_link_prediction_dataset.DistLinkPredictionDataset(*args, **kwargs)[source]#
Bases:
gigl.distributed.dist_dataset.DistDatasetThis class is inherited from GraphLearn-for-PyTorch’s DistDataset class. We override the __init__ functionality to support positive and negative edges and labels. We also override the share_ipc function to correctly serialize these new fields. We additionally introduce a build function for storing the partitioned inside of this class. We assume data in this class is only in the CPU RAM, and do not support data on GPU memory, thus simplifying the logic and tooling required compared to the base DistDataset class.
Initializes the fields of the DistDataset class. This function is called upon each serialization of the DistDataset instance. :param rank: Rank of the current process :type rank: int :param world_size: World size of the current process :type world_size: int :param edge_dir: Edge direction of the provied graph :type edge_dir: Literal[“in”, “out”]
- The below arguments are only expected to be provided when re-serializing an instance of the DistDataset class after build() has been called
graph_partition (Optional[Union[Graph, dict[EdgeType, Graph]]]): Partitioned Graph Data node_feature_partition (Optional[Union[Feature, dict[NodeType, Feature]]]): Partitioned Node Feature Data edge_feature_partition (Optional[Union[torch.Tensor, dict[EdgeType, torch.Tensor]]]): Partitioned Edge Feature Data node_labels (Optional[Union[Feature, dict[NodeType, Feature]]]): The labels of each node on the current machine node_partition_book (Optional[Union[PartitionBook, dict[NodeType, PartitionBook]]]): Node Partition Book edge_partition_book (Optional[Union[PartitionBook, dict[EdgeType, PartitionBook]]]): Edge Partition Book positive_edge_label (Optional[Union[torch.Tensor, dict[EdgeType, torch.Tensor]]]): Positive Edge Label Tensor negative_edge_label (Optional[Union[torch.Tensor, dict[EdgeType, torch.Tensor]]]): Negative Edge Label Tensor node_ids (Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]]): Node IDs on the current machine num_train: Optional[Union[int, dict[NodeType, int]]]): Number of training nodes on the current machine. Will be a dict if heterogeneous. num_val: (Optional[Union[int, dict[NodeType, int]]]): Number of validation nodes on the current machine. Will be a dict if heterogeneous. num_test: (Optional[Union[int, dict[NodeType, int]]]): Number of test nodes on the current machine. Will be a dict if heterogeneous. node_feature_info: Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]]: Dimension of node features and its data type, will be a dict if heterogeneous.
Note this will be None in the homogeneous case if the data has no node features, or will only contain node types with node features in the heterogeneous case.
- edge_feature_info: Optional[Union[FeatureInfo, dict[EdgeType, FeatureInfo]]]: Dimension of edge features and its data type, will be a dict if heterogeneous.
Note this will be None in the homogeneous case if the data has no edge features, or will only contain edge types with edge features in the heterogeneous case.