gigl.common.data.load_torch_tensors#

Attributes#

logger

Classes#

SerializedGraphMetadata

Stores information for all entities. If homogeneous, all types are of type SerializedTFRecordInfo. Otherwise, they are dictionaries with the corresponding mapping.

Functions#

`load_torch_tensors_from_tf_record`(...[, rank, ...])	Loads all torch tensors from a SerializedGraphMetadata object for all entity [node, edge, positive_label, negative_label] and edge / node types.
`remove_sampling_weight_from_edge_quantization_metadata`(...)	Remove separately stored sampling weights from edge reconstruction metadata.

Module Contents#

class gigl.common.data.load_torch_tensors.SerializedGraphMetadata[source]#

Stores information for all entities. If homogeneous, all types are of type SerializedTFRecordInfo. Otherwise, they are dictionaries with the corresponding mapping.

edge_entity_info: gigl.common.data.dataloaders.SerializedTFRecordInfo | dict[gigl.src.common.types.graph_data.EdgeType, gigl.common.data.dataloaders.SerializedTFRecordInfo][source]#

edge_quantization_metadata: gigl.types.graph.FeatureQuantizationMetadata | dict[gigl.src.common.types.graph_data.EdgeType, gigl.types.graph.FeatureQuantizationMetadata] | None = None[source]#

negative_label_entity_info: gigl.common.data.dataloaders.SerializedTFRecordInfo | dict[gigl.src.common.types.graph_data.EdgeType, gigl.common.data.dataloaders.SerializedTFRecordInfo] | None = None[source]#

node_entity_info: gigl.common.data.dataloaders.SerializedTFRecordInfo | dict[gigl.src.common.types.graph_data.NodeType, gigl.common.data.dataloaders.SerializedTFRecordInfo][source]#

node_quantization_metadata: gigl.types.graph.FeatureQuantizationMetadata | dict[gigl.src.common.types.graph_data.NodeType, gigl.types.graph.FeatureQuantizationMetadata] | None = None[source]#

positive_label_entity_info: gigl.common.data.dataloaders.SerializedTFRecordInfo | dict[gigl.src.common.types.graph_data.EdgeType, gigl.common.data.dataloaders.SerializedTFRecordInfo] | None = None[source]#

gigl.common.data.load_torch_tensors.load_torch_tensors_from_tf_record(tf_record_dataloader, serialized_graph_metadata, should_load_tensors_in_parallel, rank=0, node_tf_dataset_options=TFDatasetOptions(), edge_tf_dataset_options=TFDatasetOptions(), weight_edge_feat_name=None)[source]#

Loads all torch tensors from a SerializedGraphMetadata object for all entity [node, edge, positive_label, negative_label] and edge / node types.

Running these processes in parallel slows the runtime of each individual process, but may still result in a net speedup across all entity types. As a result, there is a tradeoff that needs to be made between parallel and sequential tensor loading, which is why we don’t parallelize across node and edge types. We enable the should_load_tensors_in_parallel to allow some customization for loading strategies based on the input data.

If RPC is initialized, node and edge feature columns are losslessly compressed, each right after its own loading process joins, unless the GIGL_FEATURE_COMPRESSION environment variable is off (see gigl.common.utils.feature_quantization.lossless_compression). When loading in parallel, node compression runs while the edge loading process is still reading. The sampling-weight edge feature is removed from the edge features before they are compressed, so it is never compressed.

With compression on and RPC initialized, every rank must call this function, with the same should_load_tensors_in_parallel, because compression joins a collective RPC all_gather. Without RPC, features are returned uncompressed.

Parameters:

tf_record_dataloader (TFRecordDataLoader) – TFRecordDataloader used for loading tensors from serialized tfrecords
serialized_graph_metadata (SerializedGraphMetadata) – Serialized graph metadata contained serialized information for loading tfrecords across node and edge types
should_load_tensors_in_parallel (bool) – Whether tensors should be loaded from serialized information in parallel or in sequence across the [node, edge, pos_label, neg_label] entity types.
rank (int) – Rank on current machine
node_tf_dataset_options (TFDatasetOptions) – The options to use for nodes when building the dataset.
edge_tf_dataset_options (TFDatasetOptions) – The options to use for edges when building the dataset.
weight_edge_feat_name (Optional[Union[str, dict[EdgeType, str]]]) – Name of the edge feature column to extract as sampling weights. The column is removed from the edge feature matrix and returned separately via LoadedGraphTensors.edge_weights. Supply a single string for homogeneous graphs or a per-edge-type dict for heterogeneous graphs.

Returns:

Unpartitioned Graph Tensors, with the quantization metadata that: describes their packed features.

Return type:

loaded_graph_tensors (LoadedGraphTensors)

Raises:

ValueError – If a loading process failed, or if GIGL_FEATURE_COMPRESSION holds an unknown value while RPC is initialized.

gigl.common.data.load_torch_tensors.remove_sampling_weight_from_edge_quantization_metadata(serialized_graph_metadata, weight_edge_feat_name)[source]#

Remove separately stored sampling weights from edge reconstruction metadata.

TFRecord loading removes the sampling-weight column from raw edge features before registering it with the weighted sampler. The resulting metadata must describe the remaining model features so batch reconstruction scatters raw and dequantized columns into the correct positions.

Parameters:

serialized_graph_metadata (SerializedGraphMetadata) – Serialized edge schema and quantization metadata.
weight_edge_feat_name (Optional[Union[str, dict[gigl.src.common.types.graph_data.EdgeType, str]]]) – Raw scalar feature configured as sampling weights.

Returns:

Quantization metadata for the model-facing edge features.

Return type:

Optional[Union[gigl.types.graph.FeatureQuantizationMetadata, dict[gigl.src.common.types.graph_data.EdgeType, gigl.types.graph.FeatureQuantizationMetadata]]]

gigl.common.data.load_torch_tensors.logger[source]#