[Example] Polish GraphSAGE link prediction example and move it to sampling folder. (#6163)

Co-authored-by: Hongzhi (Steve), Chen <chenhongzhi.nkcs@gmail.com>

examples/sampling/link_prediction.py

+"""
+This script trains and tests a GraphSAGE model for link prediction on
+large graphs using efficient and tailor-made neighbor sampling.
+
+Paper: [Inductive Representation Learning on Large Graphs]
+(https://arxiv.org/abs/1706.02216)
+
+While node classification predicts labels for nodes based on their
+local neighborhoods, link prediction assesses the likelihood of an edge
+existing between two nodes, necessitating different sampling strategies
+that account for pairs of nodes and their joint neighborhoods.
+
+Before reading this example, please familiar yourself with graphsage node
+classification by reading the example in the
+`examples/core/graphsage/link_prediction.py`
+
+If you want to train graphsage on a large graph in a distributed fashion, read
+the example in the `examples/distributed/graphsage/`.
+
+This flowchart describes the main functional sequence of the provided example.
+main
+│
+├───> Load and preprocess dataset
+│
+├───> Instantiate SAGE model
+│
+├───> train
+│     │
+│     ├───> NeighborSampler (HIGHLIGHT)
+│     │
+│     └───> Training loop
+│           │
+│           └───> SAGE.forward
+│
+└───> evaluate
+      │
+      └───> SAGE.inference
+            │
+            └───> MultiLayerFullNeighborSampler (HIGHLIGHT)
+"""
+
+import argparse
+
+import dgl
+import dgl.nn as dglnn
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import tqdm
+from dgl.dataloading import (
+    as_edge_prediction_sampler,
+    DataLoader,
+    MultiLayerFullNeighborSampler,
+    negative_sampler,
+    NeighborSampler,
+)
+from ogb.linkproppred import DglLinkPropPredDataset, Evaluator
+
+
+def to_bidirected_with_reverse_mapping(g):
+    """Convert the graph to bidirectional and return the reverse mapping.
+
+    This function transforms the input graph into its bidirectional form. It
+    then returns the newly formed bidirectional graph and the mapping that
+    represents the reverse edges. The function does not work with graphs that
+    have self-loops.
+
+    Parameters:
+    ----------
+    g : DGLGraph
+        Input graph.
+
+    Returns:
+    -------
+    DGLGraph :
+        Bidirectional graph.
+    Tensor :
+        Mapping to reverse edges.
+    """
+    # First, add reverse edges to the graph, effectively making it
+    # bidirectional. Then, simplify the resulting graph by merging any duplicate
+    # edges. The resulting simplified graph is stored in `g_simple`, and
+    # `mapping` provides information on how edges in `g_simple` correspond to
+    # edges in the original graph.
+    g_simple, mapping = dgl.to_simple(
+        dgl.add_reverse_edges(g), return_counts="count", writeback_mapping=True
+    )
+
+    # The `return_counts` option in `dgl.to_simple` returns the count of how
+    # many times each edge in the simplified graph corresponds to an edge in the
+    # original graph. This count is saved in the edge data of the returned
+    # graph with the key "count".
+    c = g_simple.edata["count"]
+    num_edges = g.num_edges()
+
+    # `mapping_offset` is an auxiliary tensor used to understand how edges in
+    # the simplified bidirectional graph (g_simple) relate to the edges in the
+    # original graph.
+    mapping_offset = torch.zeros(
+        g_simple.num_edges() + 1, dtype=g_simple.idtype
+    )
+
+    # Calculate the cumulative sum of counts to determine boundaries for each
+    # unique edge.
+    mapping_offset[1:] = c.cumsum(0)
+
+    # Sort the mapping tensor to group the same edge indices.
+    idx = mapping.argsort()
+
+    # Using the previously computed `mapping_offset`, it extracts the first
+    # index of each group, which represents the unique edge indices from the
+    # sorted mapping.
+    idx_uniq = idx[mapping_offset[:-1]]
+
+    # If an edge index is greater than or equal to the number of edges in the
+    # original graph, it indicates that this edge is a reversed edge, and the
+    # original edge index for it is (idx_uniq - num_edges). Otherwise, its
+    # reverse edge index is (idx_uniq + num_edges).
+    reverse_idx = torch.where(
+        idx_uniq >= num_edges, idx_uniq - num_edges, idx_uniq + num_edges
+    )
+    reverse_mapping = mapping[reverse_idx]
+
+    # Sanity check to ensure valid mapping.
+    src1, dst1 = g_simple.edges()
+    src2, dst2 = g_simple.find_edges(reverse_mapping)
+    assert torch.equal(src1, dst2)
+    assert torch.equal(src2, dst1)
+    return g_simple, reverse_mapping
+
+
+class SAGE(nn.Module):
+    def __init__(self, in_size, hidden_size):
+        super().__init__()
+        self.layers = nn.ModuleList()
+        # Three-layer GraphSAGE-mean.
+        self.layers.append(dglnn.SAGEConv(in_size, hidden_size, "mean"))
+        self.layers.append(dglnn.SAGEConv(hidden_size, hidden_size, "mean"))
+        self.layers.append(dglnn.SAGEConv(hidden_size, hidden_size, "mean"))
+        self.hidden_size = hidden_size
+        self.predictor = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.ReLU(),
+            nn.Linear(hidden_size, hidden_size),
+            nn.ReLU(),
+            nn.Linear(hidden_size, 1),
+        )
+
+    def forward(self, pair_graph, neg_pair_graph, blocks, x):
+        hidden_x = x
+        for layer_idx, (layer, block) in enumerate(zip(self.layers, blocks)):
+            hidden_x = layer(block, hidden_x)
+            is_last_layer = layer_idx == len(self.layers) - 1
+            if not is_last_layer:
+                hidden_x = F.relu(hidden_x)
+        pos_src, pos_dst = pair_graph.edges()
+        neg_src, neg_dst = neg_pair_graph.edges()
+        hidden_pos = self.predictor(hidden_x[pos_src] * hidden_x[pos_dst])
+        hidden_neg = self.predictor(hidden_x[neg_src] * hidden_x[neg_dst])
+        return hidden_pos, hidden_neg
+
+    def inference(self, g, device, batch_size):
+        """Layer-wise inference algorithm to compute GNN node embeddings."""
+        feat = g.ndata["feat"]
+        #####################################################################
+        # (HIGHLIGHT) Creating a MultiLayerFullNeighborSampler instance.
+        # This sampler is used in the Graph Neural Networks (GNN) training
+        # process to provide neighbor sampling, which is crucial for
+        # efficient training of GNN on large graphs.
+        #
+        # The first argument '1' indicates the number of layers for
+        # the neighbor sampling. In this case, it's set to 1, meaning
+        # only the direct neighbors of each node will be included in the
+        # sampling.
+        #
+        # The 'prefetch_node_feats' parameter specifies the node features
+        # that need to be pre-fetched during sampling. In this case, the
+        # feature named 'feat' will be pre-fetched.
+        #
+        # `prefetch` in DGL initiates data fetching operations in parallel
+        # with model computations. This ensures data is ready when the
+        # computation needs it, thereby eliminating waiting times between
+        # fetching and computing steps and reducing the I/O overhead during
+        # the training process.
+        #
+        # The difference between whether to use prefetch or not is shown:
+        #
+        # Without Prefetch:
+        # Fetch1 ──> Compute1 ──> Fetch2 ──> Compute2 ──> Fetch3 ──> Compute3
+        #
+        # With Prefetch:
+        # Fetch1 ──> Fetch2 ──> Fetch3
+        #    │          │          │
+        #    └─Compute1 └─Compute2 └─Compute3
+        #####################################################################
+        sampler = MultiLayerFullNeighborSampler(1, prefetch_node_feats=["feat"])
+        dataloader = DataLoader(
+            g,
+            torch.arange(g.num_nodes()).to(g.device),
+            sampler,
+            device=device,
+            batch_size=batch_size,
+            shuffle=False,
+            drop_last=False,
+            num_workers=0,
+        )
+        buffer_device = torch.device("cpu")
+        # Enable pin_memory for faster CPU to GPU data transfer if the model is
+        # running on a GPU.
+        pin_memory = buffer_device != device
+        for layer_idx, layer in enumerate(self.layers):
+            is_last_layer = layer_idx == len(self.layers) - 1
+            y = torch.empty(
+                g.num_nodes(),
+                self.hidden_size,
+                device=buffer_device,
+                pin_memory=pin_memory,
+            )
+            feat = feat.to(device)
+            for input_nodes, output_nodes, blocks in tqdm.tqdm(
+                dataloader, desc="Inference"
+            ):
+                x = feat[input_nodes]
+                hidden_x = layer(blocks[0], x)
+                if not is_last_layer:
+                    hidden_x = F.relu(hidden_x)
+                y[output_nodes] = hidden_x.to(buffer_device)
+            feat = y
+        return y
+
+
+@torch.no_grad()
+def compute_mrr(
+    model, evaluator, node_emb, src, dst, neg_dst, device, batch_size=500
+):
+    """Compute the Mean Reciprocal Rank (MRR) for given source and destination
+    nodes.
+
+    This function computes the MRR for a set of node pairs, dividing the task
+    into batches to handle potentially large graphs.
+    """
+    rr = torch.zeros(src.shape[0])
+    # Loop over node pairs in batches.
+    for start in tqdm.trange(0, src.shape[0], batch_size, desc="Evaluate"):
+        end = min(start + batch_size, src.shape[0])
+
+        # Concatenate positive and negative destination nodes.
+        all_dst = torch.cat([dst[start:end, None], neg_dst[start:end]], 1)
+
+        # Fetch embeddings for current batch of source and destination nodes.
+        h_src = node_emb[src[start:end]][:, None, :].to(device)
+        h_dst = node_emb[all_dst.view(-1)].view(*all_dst.shape, -1).to(device)
+
+        # Compute prediction scores using the model.
+        pred = model.predictor(h_src * h_dst).squeeze(-1)
+
+        # Evaluate the predictions to obtain MRR values.
+        input_dict = {"y_pred_pos": pred[:, 0], "y_pred_neg": pred[:, 1:]}
+        rr[start:end] = evaluator.eval(input_dict)["mrr_list"]
+    return rr.mean()
+
+
+@torch.no_grad()
+def evaluate(device, graph, edge_split, model, batch_size):
+    """Evaluate the model on validation and test sets."""
+    model.eval()
+    evaluator = Evaluator(name="ogbl-citation2")
+
+    # Compute node embeddings for the entire graph.
+    node_emb = model.inference(graph, device, batch_size)
+    results = []
+
+    # Loop over both validation and test sets.
+    for split in ["valid", "test"]:
+        src = edge_split[split]["source_node"].to(node_emb.device)
+        dst = edge_split[split]["target_node"].to(node_emb.device)
+        neg_dst = edge_split[split]["target_node_neg"].to(node_emb.device)
+
+        # Compute MRR values for the current split.
+        results.append(
+            compute_mrr(model, evaluator, node_emb, src, dst, neg_dst, device)
+        )
+    return results
+
+
+def train(args, device, g, reverse_eids, seed_edges, model, use_uva):
+    #####################################################################
+    # (HIGHLIGHT) Instantiate a NeighborSampler object for efficient
+    # training of Graph Neural Networks (GNNs) on large-scale graphs.
+    #
+    # The argument [15, 10, 5] sets the number of neighbors (fanout)
+    # to be sampled at each layer. Here, we have three layers, and
+    # 15/10/5 neighbors will be randomly selected for each node at each
+    # layer.
+    #
+    # The 'prefetch_node_feats' parameter specify the node features that
+    # needs to be pre-fetched during sampling. More details about
+    # `prefetch` can be found in the `SAGE.inference` function.
+    #
+    # (HIGHLIGHT) Modify the NeighborSampler for Edge Prediction
+    #
+    # This `as_edge_prediction_sampler` augments the original NeighborSampler
+    # to specifically handle edge prediction tasks, where not only the
+    # structure but also the relationships between nodes (edges) are of
+    # importance.
+    #
+    # - `exclude="reverse_id"` ensures that the edges corresponding to the
+    #   reverse of the original edges are excluded during sampling, given that
+    #   reverse edges can introduce unnecessary redundancy in edge prediction.
+    #
+    # - `reverse_eids=reverse_eids` specifies the IDs of the reverse edges.
+    #   This information is vital so the sampler knows which edges to avoid.
+    #
+    # - The negative sampling strategy is specified using the
+    #   `negative_sampler`. Here, a uniform negative sampling method is
+    #   employed, where a negative sample (an edge that doesn't exist in the
+    #   original graph) is uniformly drawn from the set of all possible edges.
+    #
+    # The modified sampler is tailor-made for scenarios where the goal is
+    # not just to learn node representations, but also to predict the
+    # likelihood of an edge existing between two nodes (link prediction).
+    #####################################################################
+    sampler = NeighborSampler([15, 10, 5], prefetch_node_feats=["feat"])
+    sampler = as_edge_prediction_sampler(
+        sampler,
+        exclude="reverse_id",
+        reverse_eids=reverse_eids,
+        negative_sampler=negative_sampler.Uniform(1),
+    )
+
+    dataloader = DataLoader(
+        g,
+        seed_edges,
+        sampler,
+        device=device,
+        batch_size=args.batch_size,
+        shuffle=True,
+        drop_last=False,
+        # If `g` is on gpu or `use_uva` is True, `num_workers` must be zero,
+        # otherwise it will cause error.
+        num_workers=0,
+        use_uva=use_uva,
+    )
+    opt = torch.optim.Adam(model.parameters(), lr=args.lr)
+    for epoch in range(10):
+        model.train()
+        total_loss = 0
+        # A block is a graph consisting of two sets of nodes: the
+        # source nodes and destination nodes. The source and destination
+        # nodes can have multiple node types. All the edges connect from
+        # source nodes to destination nodes.
+        # For more details: https://discuss.dgl.ai/t/what-is-the-block/2932.
+        for it, (input_nodes, pair_graph, neg_pair_graph, blocks) in enumerate(
+            dataloader
+        ):
+            # The input features from the source nodes in the first layer's
+            # computation graph.
+            x = blocks[0].srcdata["feat"]
+            pos_score, neg_score = model(pair_graph, neg_pair_graph, blocks, x)
+            score = torch.cat([pos_score, neg_score])
+
+            # Create true labels for positive and negative samples.
+            pos_label = torch.ones_like(pos_score)
+            neg_label = torch.zeros_like(neg_score)
+            labels = torch.cat([pos_label, neg_label])
+
+            # Compute the binary cross-entropy loss.
+            loss = F.binary_cross_entropy_with_logits(score, labels)
+            opt.zero_grad()
+            loss.backward()
+            opt.step()
+            total_loss += loss.item()
+        print(f"Epoch {epoch:05d} | Loss {total_loss / (it + 1):.4f}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--lr",
+        type=float,
+        default=0.0005,
+        help="Learning rate. Default: 0.0005",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=512,
+        help="Batch size. Default: 512",
+    )
+    parser.add_argument(
+        "--mode",
+        default="mixed",
+        choices=["cpu", "mixed", "puregpu"],
+        help="Training mode. 'cpu' for CPU training, 'mixed' for CPU-GPU mixed "
+        "training, 'puregpu' for pure-GPU training.",
+    )
+    return parser.parse_args()
+
+
+def main(args):
+    if not torch.cuda.is_available():
+        args.mode = "cpu"
+    print(f"Training in {args.mode} mode.")
+
+    # Load and preprocess dataset.
+    print("Loading data")
+    dataset = DglLinkPropPredDataset("ogbl-citation2")
+    g = dataset[0]
+    g = g.to("cuda" if args.mode == "puregpu" else "cpu")
+
+    # Whether use Unified Virtual Addressing (UVA) for CUDA computation.
+    use_uva = args.mode == "mixed"
+    device = torch.device("cpu" if args.mode == "cpu" else "cuda")
+
+    # Convert the graph to its bidirectional form.
+    g, reverse_eids = to_bidirected_with_reverse_mapping(g)
+    reverse_eids = reverse_eids.to(device)
+    seed_edges = torch.arange(g.num_edges()).to(device)
+    edge_split = dataset.get_edge_split()
+
+    # Create GraphSAGE model.
+    in_size = g.ndata["feat"].shape[1]
+    model = SAGE(in_size, 256).to(device)
+
+    # Model training.
+    print("Training...")
+    train(args, device, g, reverse_eids, seed_edges, model, use_uva)
+
+    # Validate/Test the model.
+    print("Validation/Testing...")
+    valid_mrr, test_mrr = evaluate(
+        device, g, edge_split, model, batch_size=1000
+    )
+    print(
+        f"Validation MRR {valid_mrr.item():.4f}, Test MRR {test_mrr.item():.4f}"
+    )
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)