link_prediction.py 16 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
"""
This script trains and tests a GraphSAGE model for link prediction on
large graphs using graphbolt dataloader.

Paper: [Inductive Representation Learning on Large Graphs]
(https://arxiv.org/abs/1706.02216)

Unlike previous dgl examples, we've utilized the newly defined dataloader
from GraphBolt. This example will help you grasp how to build an end-to-end
training pipeline using GraphBolt.

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.

TODO: Add the link_prediction.py example to core/graphsage.
Before reading this example, please familiar yourself with graphsage link
prediction 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

├───> OnDiskDataset pre-processing

├───> Instantiate SAGE model

├───> train
│     │
│     ├───> Get graphbolt dataloader (HIGHLIGHT)
│     │
│     └───> Training loop
│           │
│           ├───> SAGE.forward
│           │
│           └───> Validation set evaluation

└───> Test set evaluation
"""
import argparse
45
import time
46
from functools import partial
47
48
49
50
51
52
53

import dgl.graphbolt as gb
import dgl.nn as dglnn
import torch
import torch.nn as nn
import torch.nn.functional as F
import tqdm
54
from torchmetrics.retrieval import RetrievalMRR
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81


class SAGE(nn.Module):
    def __init__(self, in_size, hidden_size):
        super().__init__()
        self.layers = nn.ModuleList()
        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, 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)
        return hidden_x

82
    def inference(self, graph, features, dataloader, storage_device):
83
        """Conduct layer-wise inference to get all the node embeddings."""
84
85
        pin_memory = storage_device == "pinned"
        buffer_device = torch.device("cpu" if pin_memory else storage_device)
86
87
88
89
90
91
92
93
94
95
96
97

        print("Start node embedding inference.")
        for layer_idx, layer in enumerate(self.layers):
            is_last_layer = layer_idx == len(self.layers) - 1

            y = torch.empty(
                graph.total_num_nodes,
                self.hidden_size,
                dtype=torch.float32,
                device=buffer_device,
                pin_memory=pin_memory,
            )
98
99
100
            for data in tqdm.tqdm(dataloader):
                # len(blocks) = 1
                hidden_x = layer(data.blocks[0], data.node_features["feat"])
101
102
                if not is_last_layer:
                    hidden_x = F.relu(hidden_x)
103
                # By design, our seed nodes are contiguous.
104
                y[data.seeds[0] : data.seeds[-1] + 1] = hidden_x.to(
105
106
                    buffer_device, non_blocking=True
                )
107
108
            if not is_last_layer:
                features.update("node", None, "feat", y)
109
110
111

        return y

112
113
114
115
116

def create_dataloader(args, graph, features, itemset, is_train=True):
    """Get a GraphBolt version of a dataloader for link prediction tasks. This
    function demonstrates how to utilize functional forms of datapipes in
    GraphBolt. Alternatively, you can create a datapipe using its class
117
118
    constructor. For a more detailed tutorial, please read the examples in
    `dgl/notebooks/graphbolt/walkthrough.ipynb`.
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
    """

    ############################################################################
    # [Input]:
    # 'itemset': The current dataset.
    # 'args.batch_size': Specify the number of samples to be processed together,
    # referred to as a 'mini-batch'. (The term 'mini-batch' is used here to
    # indicate a subset of the entire dataset that is processed together. This
    # is in contrast to processing the entire dataset, known as a 'full batch'.)
    # 'is_train': Determining if data should be shuffled. (Shuffling is
    # generally used only in training to improve model generalization. It's
    # not used in validation and testing as the focus there is to evaluate
    # performance rather than to learn from the data.)
    # [Output]:
    # An ItemSampler object for handling mini-batch sampling.
    # [Role]:
    # Initialize the ItemSampler to sample mini-batche from the dataset.
    ############################################################################
    datapipe = gb.ItemSampler(
138
139
140
        itemset,
        batch_size=args.train_batch_size if is_train else args.eval_batch_size,
        shuffle=is_train,
141
142
    )

143
144
145
146
147
148
149
150
151
152
    ############################################################################
    # [Input]:
    # 'device': The device to copy the data to.
    # [Output]:
    # A CopyTo object to copy the data to the specified device. Copying here
    # ensures that the rest of the operations run on the GPU.
    ############################################################################
    if args.storage_device != "cpu":
        datapipe = datapipe.copy_to(device=args.device)

153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
    ############################################################################
    # [Input]:
    # 'args.neg_ratio': Specify the ratio of negative to positive samples.
    # (E.g., if neg_ratio is 1, for each positive sample there will be 1
    # negative sample.)
    # 'graph': The overall network topology for negative sampling.
    # [Output]:
    # A UniformNegativeSampler object that will handle the generation of
    # negative samples for link prediction tasks.
    # [Role]:
    # Initialize the UniformNegativeSampler for negative sampling in link
    # prediction.
    # [Note]:
    # If 'is_train' is False, the UniformNegativeSampler will not be used.
    # Since, in validation and testing, the itemset already contains the
    # negative edges information.
    ############################################################################
    if is_train:
171
        datapipe = datapipe.sample_uniform_negative(graph, args.neg_ratio)
172
173
174
175
176
177
178
179
180
181
182
183

    ############################################################################
    # [Input]:
    # 'datapipe' is either 'ItemSampler' or 'UniformNegativeSampler' depending
    # on whether training is needed ('is_train'),
    # 'graph': The network topology for sampling.
    # 'args.fanout': Number of neighbors to sample per node.
    # [Output]:
    # A NeighborSampler object to sample neighbors.
    # [Role]:
    # Initialize a neighbor sampler for sampling the neighborhoods of nodes.
    ############################################################################
184
185
186
    datapipe = datapipe.sample_neighbor(
        graph, args.fanout if is_train else [-1]
    )
187

188
189
190
191
192
193
194
195
196
197
198
199
200
    ############################################################################
    # [Input]:
    # 'gb.exclude_seed_edges': Function to exclude seed edges, optionally
    # including their reverse edges, from the sampled subgraphs in the
    # minibatch.
    # [Output]:
    # A MiniBatchTransformer object with excluded seed edges.
    # [Role]:
    # During the training phase of link prediction, negative edges are
    # sampled. It's essential to exclude the seed edges from the process
    # to ensure that positive samples are not inadvertently included within
    # the negative samples.
    ############################################################################
201
    if is_train and args.exclude_edges:
202
203
204
        datapipe = datapipe.transform(
            partial(gb.exclude_seed_edges, include_reverse_edges=True)
        )
205

206
207
208
209
210
211
212
213
    ############################################################################
    # [Input]:
    # 'features': The node features.
    # 'node_feature_keys': The node feature keys (list) to be fetched.
    # [Output]:
    # A FeatureFetcher object to fetch node features.
    # [Role]:
    # Initialize a feature fetcher for fetching features of the sampled
214
    # subgraphs.
215
    ############################################################################
216
    datapipe = datapipe.fetch_feature(features, node_feature_keys=["feat"])
217
218
219
220
221
222
223

    ############################################################################
    # [Input]:
    # 'device': The device to copy the data to.
    # [Output]:
    # A CopyTo object to copy the data to the specified device.
    ############################################################################
224
225
    if args.storage_device == "cpu":
        datapipe = datapipe.copy_to(device=args.device)
226
227
228
229
230
231

    ############################################################################
    # [Input]:
    # 'datapipe': The datapipe object to be used for data loading.
    # 'args.num_workers': The number of processes to be used for data loading.
    # [Output]:
232
    # A DataLoader object to handle data loading.
233
234
235
    # [Role]:
    # Initialize a multi-process dataloader to load the data in parallel.
    ############################################################################
236
    dataloader = gb.DataLoader(
237
238
239
240
241
242
243
244
245
        datapipe,
        num_workers=args.num_workers,
    )

    # Return the fully-initialized DataLoader object.
    return dataloader


@torch.no_grad()
246
def compute_mrr(args, model, node_emb, seeds, labels, indexes):
247
248
249
250
251
252
253
    """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.
    """

254
    preds = torch.empty(seeds.shape[0], device=indexes.device)
255
256
257
258
259
260
261
262
    mrr = RetrievalMRR()
    seeds_src, seeds_dst = seeds.T
    # The constant number is 1001, due to negtive ratio in the `ogbl-citation2`
    # dataset is 1000.
    eval_size = args.eval_batch_size * 1001
    # Loop over node pairs in batches.
    for start in tqdm.trange(0, seeds_src.shape[0], eval_size, desc="Evaluate"):
        end = min(start + eval_size, seeds_src.shape[0])
263
264

        # Fetch embeddings for current batch of source and destination nodes.
265
266
        h_src = node_emb[seeds_src[start:end]].to(args.device)
        h_dst = node_emb[seeds_dst[start:end]].to(args.device)
267
268

        # Compute prediction scores using the model.
269
270
271
        pred = model.predictor(h_src * h_dst).squeeze()
        preds[start:end] = pred
    return mrr(preds, labels, indexes=indexes)
272
273
274
275
276
277


@torch.no_grad()
def evaluate(args, model, graph, features, all_nodes_set, valid_set, test_set):
    """Evaluate the model on validation and test sets."""
    model.eval()
278
279

    dataloader = create_dataloader(
280
        args, graph, features, all_nodes_set, is_train=False
281
282
    )

283
    # Compute node embeddings for the entire graph.
284
    node_emb = model.inference(graph, features, dataloader, args.storage_device)
285
    results = []
286

287
288
289
    # Loop over both validation and test sets.
    for split in [valid_set, test_set]:
        # Unpack the item set.
290
291
292
        seeds = split._items[0].to(node_emb.device)
        labels = split._items[1].to(node_emb.device)
        indexes = split._items[2].to(node_emb.device)
293

294
295
        # Compute MRR values for the current split.
        results.append(
296
            compute_mrr(args, model, node_emb, seeds, labels, indexes)
297
298
        )
    return results
299
300


301
def train(args, model, graph, features, train_set):
302
303
304
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
    dataloader = create_dataloader(args, graph, features, train_set)

305
    for epoch in range(args.epochs):
306
307
        model.train()
        total_loss = 0
308
        start_epoch_time = time.time()
309
        for step, data in tqdm.tqdm(enumerate(dataloader)):
310
            # Get node pairs with labels for loss calculation.
311
312
            compacted_seeds = data.compacted_seeds.T
            labels = data.labels
313

314
            node_feature = data.node_features["feat"]
315
            blocks = data.blocks
316
317
318
319

            # Get the embeddings of the input nodes.
            y = model(blocks, node_feature)
            logits = model.predictor(
320
                y[compacted_seeds[0]] * y[compacted_seeds[1]]
321
322
323
324
325
326
327
328
329
            ).squeeze()

            # Compute loss.
            loss = F.binary_cross_entropy_with_logits(logits, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()
330
            if step + 1 == args.early_stop:
331
332
                # Early stopping requires a new dataloader to reset its state.
                dataloader = create_dataloader(args, graph, features, train_set)
333
                break
334

335
336
337
338
339
340
        end_epoch_time = time.time()
        print(
            f"Epoch {epoch:05d} | "
            f"Loss {(total_loss) / (step + 1):.4f} | "
            f"Time {(end_epoch_time - start_epoch_time):.4f} s"
        )
341
342
343
344
345
346
347


def parse_args():
    parser = argparse.ArgumentParser(description="OGBL-Citation2 (GraphBolt)")
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--lr", type=float, default=0.0005)
    parser.add_argument("--neg-ratio", type=int, default=1)
348
    parser.add_argument("--train-batch-size", type=int, default=512)
349
    parser.add_argument("--eval-batch-size", type=int, default=1024)
350
    parser.add_argument("--num-workers", type=int, default=0)
351
352
353
354
355
356
    parser.add_argument(
        "--early-stop",
        type=int,
        default=0,
        help="0 means no early stop, otherwise stop at the input-th step",
    )
357
358
359
360
361
362
    parser.add_argument(
        "--fanout",
        type=str,
        default="15,10,5",
        help="Fan-out of neighbor sampling. Default: 15,10,5",
    )
363
364
365
366
367
368
    parser.add_argument(
        "--exclude-edges",
        type=int,
        default=1,
        help="Whether to exclude reverse edges during sampling. Default: 1",
    )
369
    parser.add_argument(
370
371
372
373
374
        "--mode",
        default="pinned-cuda",
        choices=["cpu-cpu", "cpu-cuda", "pinned-cuda", "cuda-cuda"],
        help="Dataset storage placement and Train device: 'cpu' for CPU and RAM,"
        " 'pinned' for pinned memory in RAM, 'cuda' for GPU and GPU memory.",
375
376
377
378
379
380
    )
    return parser.parse_args()


def main(args):
    if not torch.cuda.is_available():
381
382
383
384
        args.mode = "cpu-cpu"
    print(f"Training in {args.mode} mode.")
    args.storage_device, args.device = args.mode.split("-")
    args.device = torch.device(args.device)
385
386
387

    # Load and preprocess dataset.
    print("Loading data")
388
    dataset = gb.BuiltinDataset("ogbl-citation2").load()
389
390

    # Move the dataset to the selected storage.
391
392
393
394
395
396
    if args.storage_device == "pinned":
        graph = dataset.graph.pin_memory_()
        features = dataset.feature.pin_memory_()
    else:
        graph = dataset.graph.to(args.storage_device)
        features = dataset.feature.to(args.storage_device)
397

398
399
400
    train_set = dataset.tasks[0].train_set
    args.fanout = list(map(int, args.fanout.split(",")))

401
    in_size = features.size("node", None, "feat")[0]
402
    hidden_channels = 256
403
404
    args.device = torch.device(args.device)
    model = SAGE(in_size, hidden_channels).to(args.device)
405
406
407

    # Model training.
    print("Training...")
408
    train(args, model, graph, features, train_set)
409

410
411
412
    # Test the model.
    print("Testing...")
    test_set = dataset.tasks[0].test_set
413
414
415
416
417
418
419
420
421
    valid_set = dataset.tasks[0].validation_set
    all_nodes_set = dataset.all_nodes_set
    valid_mrr, test_mrr = evaluate(
        args, model, graph, features, all_nodes_set, valid_set, test_set
    )
    print(
        f"Validation MRR {valid_mrr.item():.4f}, "
        f"Test MRR {test_mrr.item():.4f}"
    )
422
423
424
425
426


if __name__ == "__main__":
    args = parse_args()
    main(args)