Codex/graph store multi supervision edge types

examples/link_prediction/graph_store/heterogeneous_training.py

@@ -211,6 +211,7 @@ def _setup_dataloaders(
         channel_size=sampling_worker_shared_channel_size,
         process_start_gap_seconds=process_start_gap_seconds,
         shuffle=shuffle,
+        use_edge_index_output=True,
     )
 
     print(f"---Rank {rank} finished setting up main loader for split={split}")
@@ -294,23 +295,16 @@ def _compute_loss(
     )
 
     # Extracting local query, random negative, positive, hard_negative, and random_negative indices.
-    query_node_idx: torch.Tensor = torch.arange(
-        main_data[query_node_type].batch_size
-    ).to(device)
+    query_node_idx = torch.arange(main_data[query_node_type].batch_size, device=device)
     random_negative_batch_size = random_negative_data[labeled_node_type].batch_size
 
-    positive_idx: torch.Tensor = torch.cat(list(main_data.y_positive.values())).to(
-        device
-    )
-    repeated_query_node_idx = query_node_idx.repeat_interleave(
-        torch.tensor([len(v) for v in main_data.y_positive.values()]).to(device)
-    )
+    positive_label_edge_index: torch.Tensor = main_data.y_positive
+    repeated_query_node_idx = query_node_idx[positive_label_edge_index[0]]
+    positive_idx = positive_label_edge_index[1]
     if hasattr(main_data, "y_negative"):
-        hard_negative_idx: torch.Tensor = torch.cat(
-            list(main_data.y_negative.values())
-        ).to(device)
+        hard_negative_idx: torch.Tensor = main_data.y_negative[1]
     else:
-        hard_negative_idx = torch.empty(0, dtype=torch.long).to(device)
+        hard_negative_idx = torch.empty(0, dtype=torch.long, device=device)
 
     # Use local IDs to get the corresponding embeddings in the tensors
 

examples/link_prediction/graph_store/homogeneous_training.py

@@ -241,6 +241,7 @@ def _setup_dataloaders(
         channel_size=sampling_worker_shared_channel_size,
         process_start_gap_seconds=process_start_gap_seconds,
         shuffle=shuffle,
+        use_edge_index_output=True,
     )
 
     logger.info(f"---Rank {rank} finished setting up main loader for split={split}")
@@ -302,21 +303,16 @@ def _compute_loss(
     main_embeddings = model(data=main_data, device=device)
     random_negative_embeddings = model(data=random_negative_data, device=device)
 
-    query_node_idx: torch.Tensor = torch.arange(main_data.batch_size).to(device)
+    query_node_idx = torch.arange(main_data.batch_size, device=device)
     random_negative_batch_size = random_negative_data.batch_size
 
-    positive_idx: torch.Tensor = torch.cat(list(main_data.y_positive.values())).to(
-        device
-    )
-    repeated_query_node_idx = query_node_idx.repeat_interleave(
-        torch.tensor([len(v) for v in main_data.y_positive.values()]).to(device)
-    )
+    positive_label_edge_index: torch.Tensor = main_data.y_positive
+    repeated_query_node_idx = query_node_idx[positive_label_edge_index[0]]
+    positive_idx = positive_label_edge_index[1]
     if hasattr(main_data, "y_negative"):
-        hard_negative_idx: torch.Tensor = torch.cat(
-            list(main_data.y_negative.values())
-        ).to(device)
+        hard_negative_idx: torch.Tensor = main_data.y_negative[1]
     else:
-        hard_negative_idx = torch.empty(0, dtype=torch.long).to(device)
+        hard_negative_idx = torch.empty(0, dtype=torch.long, device=device)
 
     # Use local IDs to get the corresponding embeddings in the tensors
 

examples/link_prediction/heterogeneous_training.py

@@ -144,6 +144,7 @@ def _setup_dataloaders(
         # This is done so that each process on the current machine which initializes a `main_loader` doesn't compete for memory, causing potential OOM
         process_start_gap_seconds=process_start_gap_seconds,
         shuffle=shuffle,
+        use_edge_index_output=True,
     )
 
     logger.info(f"---Rank {rank} finished setting up main loader")
@@ -218,25 +219,16 @@ def _compute_loss(
     # Local in this case refers to the local index in the batch, while global subsequently refers to the node's unique global ID across all nodes in the dataset.
     # Global ids are stored in data.node, ex. `data.node = [50, 20, 10]` where the `0` is the local index for global id `50`. Note that each global id in data.node is unique.
     # TODO (mkolodner-sc): If GLT migrates towards using PyG's `.n_id` field instead of `.node`, update this code.
-    query_node_idx: torch.Tensor = torch.arange(
-        main_data[query_node_type].batch_size
-    ).to(device)
+    query_node_idx = torch.arange(main_data[query_node_type].batch_size, device=device)
     random_negative_batch_size = random_negative_data[labeled_node_type].batch_size
 
-    # main_data.y_positive is a dict[query_node_local_index: int, labeled_node_local_indices: torch.Tensor], even in the heterogeneous setting.
-    positive_idx: torch.Tensor = torch.cat(list(main_data.y_positive.values())).to(
-        device
-    )
-    # We also extract a repeated query node index tensor which upsamples each query node based on the number of positives it has
-    repeated_query_node_idx = query_node_idx.repeat_interleave(
-        torch.tensor([len(v) for v in main_data.y_positive.values()]).to(device)
-    )
+    positive_label_edge_index: torch.Tensor = main_data.y_positive
+    repeated_query_node_idx = query_node_idx[positive_label_edge_index[0]]
+    positive_idx = positive_label_edge_index[1]
     if hasattr(main_data, "y_negative"):
-        hard_negative_idx: torch.Tensor = torch.cat(
-            list(main_data.y_negative.values())
-        ).to(device)
+        hard_negative_idx: torch.Tensor = main_data.y_negative[1]
     else:
-        hard_negative_idx = torch.empty(0, dtype=torch.long).to(device)
+        hard_negative_idx = torch.empty(0, dtype=torch.long, device=device)
 
     # Use local IDs to get the corresponding embeddings in the tensors
 

examples/link_prediction/homogeneous_training.py

@@ -134,6 +134,7 @@ def _setup_dataloaders(
         # This is done so that each process on the current machine which initializes a `main_loader` doesn't compete for memory, causing potential OOM
         process_start_gap_seconds=process_start_gap_seconds,
         shuffle=shuffle,
+        use_edge_index_output=True,
     )
 
     logger.info(f"---Rank {rank} finished setting up main loader")
@@ -187,23 +188,17 @@ def _compute_loss(
     # Extracting local query, random negative, positive, hard_negative, and random_negative indices.
     # Local in this case refers to the local index in the batch, while global subsequently refers to the node's unique global ID across all nodes in the dataset.
     # Global ids are stored in data.node, ex. `data.node = [50, 20, 10]` where the `0` is the local index for global id `50`. Note that each global id in data.node is unique.
-    query_node_idx: torch.Tensor = torch.arange(main_data.batch_size).to(device)
+    query_node_idx = torch.arange(main_data.batch_size, device=device)
     random_negative_batch_size = random_negative_data.batch_size
 
-    # main_data.y_positive is a dict[query_node_local_index: int, labeled_node_local_indices: torch.Tensor]
-    positive_idx: torch.Tensor = torch.cat(list(main_data.y_positive.values())).to(
-        device
-    )
-    # We also extract a repeated query node index tensor which upsamples each query node based on the number of positives it has
-    repeated_query_node_idx = query_node_idx.repeat_interleave(
-        torch.tensor([len(v) for v in main_data.y_positive.values()]).to(device)
-    )
+    # Label edge indices are [anchor_local_id, label_local_id] pairs.
+    positive_label_edge_index: torch.Tensor = main_data.y_positive
+    repeated_query_node_idx = query_node_idx[positive_label_edge_index[0]]
+    positive_idx = positive_label_edge_index[1]
     if hasattr(main_data, "y_negative"):
-        hard_negative_idx: torch.Tensor = torch.cat(
-            list(main_data.y_negative.values())
-        ).to(device)
+        hard_negative_idx: torch.Tensor = main_data.y_negative[1]
     else:
-        hard_negative_idx = torch.empty(0, dtype=torch.long).to(device)
+        hard_negative_idx = torch.empty(0, dtype=torch.long, device=device)
 
     # Use local IDs to get the corresponding embeddings in the tensors