Link Prediction¶
This tutorial will show how to train a multi-layer GraphSAGE for link prediction on CoraGraphDataset. The dataset contains 2708 nodes and 10556 edges.
By the end of this tutorial, you will be able to
Train a GNN model for link prediction on target device with DGL’s neighbor sampling components.
Install DGL package¶
[1]:
# Install required packages.
import os
import torch
import numpy as np
os.environ['TORCH'] = torch.__version__
os.environ['DGLBACKEND'] = "pytorch"
# Install the CPU version. If you want to install CUDA version, please
# refer to https://www.dgl.ai/pages/start.html.
device = torch.device("cpu")
!pip install --pre dgl -f https://data.dgl.ai/wheels-test/repo.html
try:
import dgl
import dgl.graphbolt as gb
installed = True
except ImportError as error:
installed = False
print(error)
print("DGL installed!" if installed else "DGL not found!")
Looking in links: https://data.dgl.ai/wheels-test/repo.html
Requirement already satisfied: dgl in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages/dgl-2.0.0-py3.7-linux-x86_64.egg (2.0.0)
Requirement already satisfied: numpy>=1.14.0 in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from dgl) (1.21.6)
Requirement already satisfied: scipy>=1.1.0 in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from dgl) (1.7.3)
Requirement already satisfied: networkx>=2.1 in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from dgl) (2.6.3)
Requirement already satisfied: requests>=2.19.0 in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from dgl) (2.31.0)
Requirement already satisfied: tqdm in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from dgl) (4.66.1)
Requirement already satisfied: psutil>=5.8.0 in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from dgl) (5.9.7)
Requirement already satisfied: torchdata>=0.5.0 in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from dgl) (0.5.1)
Requirement already satisfied: charset-normalizer<4,>=2 in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from requests>=2.19.0->dgl) (3.3.2)
Requirement already satisfied: idna<4,>=2.5 in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from requests>=2.19.0->dgl) (3.6)
Requirement already satisfied: urllib3<3,>=1.21.1 in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from requests>=2.19.0->dgl) (2.0.7)
Requirement already satisfied: certifi>=2017.4.17 in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from requests>=2.19.0->dgl) (2023.11.17)
Requirement already satisfied: portalocker>=2.0.0 in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from torchdata>=0.5.0->dgl) (2.7.0)
Requirement already satisfied: torch==1.13.1 in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from torchdata>=0.5.0->dgl) (1.13.1)
Requirement already satisfied: typing-extensions in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from torch==1.13.1->torchdata>=0.5.0->dgl) (4.7.1)
Requirement already satisfied: nvidia-cuda-runtime-cu11==11.7.99 in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from torch==1.13.1->torchdata>=0.5.0->dgl) (11.7.99)
Requirement already satisfied: nvidia-cudnn-cu11==8.5.0.96 in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from torch==1.13.1->torchdata>=0.5.0->dgl) (8.5.0.96)
Requirement already satisfied: nvidia-cublas-cu11==11.10.3.66 in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from torch==1.13.1->torchdata>=0.5.0->dgl) (11.10.3.66)
Requirement already satisfied: nvidia-cuda-nvrtc-cu11==11.7.99 in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from torch==1.13.1->torchdata>=0.5.0->dgl) (11.7.99)
Requirement already satisfied: setuptools in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch==1.13.1->torchdata>=0.5.0->dgl) (68.0.0)
Requirement already satisfied: wheel in /home/ubuntu/prod-doc/readthedocs.org/user_builds/dgl/envs/2.0.x/lib/python3.7/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch==1.13.1->torchdata>=0.5.0->dgl) (0.42.0)
WARNING:root:The OGB package is out of date. Your version is 1.2.4, while the latest version is 1.3.6.
DGL installed!
Loading Dataset¶
cora is already prepared as BuiltinDataset in GraphBolt.
[2]:
dataset = gb.BuiltinDataset("cora").load()
Downloading datasets/cora.zip from https://data.dgl.ai/dataset/graphbolt/cora.zip...
datasets/cora.zip: 100%|██████████| 240k/240k [00:00<00:00, 11.3MB/s]
Extracting file to datasets
Start to preprocess the on-disk dataset.
Finish preprocessing the on-disk dataset.
Dataset consists of graph, feature and tasks. You can get the training-validation-test set from the tasks. Seed nodes and corresponding labels are already stored in each training-validation-test set. This dataset contains 2 tasks, one for node classification and the other for link prediction. We will use the link prediction task.
[3]:
graph = dataset.graph
feature = dataset.feature
train_set = dataset.tasks[1].train_set
test_set = dataset.tasks[1].test_set
task_name = dataset.tasks[1].metadata["name"]
print(f"Task: {task_name}.")
Task: link_prediction.
Defining Neighbor Sampler and Data Loader in DGL¶
Different from the link prediction tutorial for full graph, a common practice to train GNN on large graphs is to iterate over the edges in minibatches, since computing the probability of all edges is usually impossible. For each minibatch of edges, you compute the output representation of their incident nodes using neighbor sampling and GNN, in a similar fashion introduced in the node classification tutorial.
To perform link prediction, you need to specify a negative sampler. DGL provides builtin negative samplers such as dgl.graphbolt.UniformNegativeSampler. Here this tutorial uniformly draws 5 negative examples per positive example.
Except for the negative sampler, the rest of the code is identical to the node classification tutorial.
[4]:
from functools import partial
datapipe = gb.ItemSampler(train_set, batch_size=256, shuffle=True)
datapipe = datapipe.sample_uniform_negative(graph, 5)
datapipe = datapipe.sample_neighbor(graph, [5, 5, 5])
datapipe = datapipe.transform(partial(gb.exclude_seed_edges, include_reverse_edges=True))
datapipe = datapipe.fetch_feature(feature, node_feature_keys=["feat"])
datapipe = datapipe.copy_to(device)
train_dataloader = gb.DataLoader(datapipe)
You can peek one minibatch from train_dataloader and see what it will give you.
[5]:
data = next(iter(train_dataloader))
print(f"MiniBatch: {data}")
MiniBatch: MiniBatch(seed_nodes=None,
sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([ 0, 2, 5, ..., 7687, 7689, 7690]),
indices=tensor([1311, 658, 445, ..., 2549, 2257, 2269]),
),
original_row_node_ids=tensor([1970, 1536, 875, ..., 1166, 1678, 1364]),
original_edge_ids=None,
original_column_node_ids=tensor([1970, 1536, 875, ..., 1392, 881, 2527]),
),
SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([ 0, 2, 5, ..., 7100, 7103, 7106]),
indices=tensor([1311, 658, 445, ..., 2280, 2061, 1309]),
),
original_row_node_ids=tensor([1970, 1536, 875, ..., 1392, 881, 2527]),
original_edge_ids=None,
original_column_node_ids=tensor([1970, 1536, 875, ..., 762, 1047, 1588]),
),
SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([ 0, 2, 5, ..., 3872, 3875, 3876]),
indices=tensor([1311, 658, 445, ..., 2280, 2281, 1940]),
),
original_row_node_ids=tensor([1970, 1536, 875, ..., 762, 1047, 1588]),
original_edge_ids=None,
original_column_node_ids=tensor([1970, 1536, 875, ..., 1695, 2333, 2620]),
)],
positive_node_pairs=(tensor([ 0, 1, 2, 3, 4, 5, 6, 7, 670, 8, 9, 10, 11, 12,
13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 767, 25,
26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53,
54, 55, 56, 57, 58, 59, 17, 60, 61, 62, 63, 64, 65, 66,
67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 334, 77, 78, 79,
80, 721, 81, 658, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91,
658, 92, 7, 93, 94, 95, 96, 658, 97, 98, 66, 99, 100, 101,
730, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114,
115, 116, 115, 333, 117, 118, 119, 120, 63, 121, 122, 845, 123, 333,
124, 17, 653, 125, 88, 752, 126, 127, 128, 129, 109, 130, 131, 132,
133, 45, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145,
22, 146, 147, 148, 867, 149, 150, 151, 813, 152, 153, 154, 155, 156,
101, 88, 694, 157, 158, 159, 778, 127, 160, 161, 162, 163, 164, 133,
165, 166, 789, 861, 167, 168, 169, 137, 170, 171, 172, 173, 174, 175,
176, 177, 689, 178, 379, 179, 180, 181, 182, 183, 184, 801, 185, 186,
187, 188, 189, 190, 191, 192, 826, 193, 194, 195, 196, 643, 788, 197,
198, 199, 940, 946, 200, 201, 202, 203, 952, 6, 795, 204, 205, 454,
911, 206, 207, 427]),
tensor([ 208, 209, 210, 211, 163, 212, 213, 214, 215, 216, 217, 218,
219, 20, 220, 221, 222, 348, 223, 224, 225, 226, 227, 228,
229, 230, 231, 232, 466, 233, 234, 235, 236, 237, 238, 239,
178, 680, 240, 241, 16, 242, 440, 178, 243, 1007, 244, 245,
36, 246, 247, 248, 993, 249, 250, 251, 252, 253, 254, 255,
256, 257, 258, 259, 960, 260, 989, 261, 262, 263, 772, 353,
264, 768, 265, 266, 131, 267, 17, 491, 668, 7, 268, 767,
269, 270, 271, 829, 272, 273, 1040, 1066, 487, 274, 768, 275,
334, 812, 276, 62, 448, 694, 277, 278, 279, 503, 213, 280,
124, 281, 77, 701, 282, 758, 283, 284, 118, 122, 112, 285,
81, 787, 286, 287, 88, 514, 908, 824, 288, 289, 209, 936,
290, 291, 292, 933, 293, 294, 295, 296, 377, 297, 260, 298,
299, 300, 797, 301, 302, 303, 854, 304, 305, 262, 705, 275,
439, 1129, 306, 1097, 302, 307, 308, 309, 694, 155, 310, 47,
311, 694, 197, 400, 312, 313, 51, 694, 314, 315, 316, 317,
318, 1066, 319, 113, 320, 16, 7, 321, 322, 26, 323, 324,
694, 325, 326, 327, 328, 329, 330, 331, 809, 332, 333, 761,
334, 4, 335, 336, 337, 338, 339, 340, 764, 341, 342, 343,
344, 345, 346, 347, 348, 17, 333, 349, 350, 351, 352, 353,
354, 355, 79, 785, 814, 356, 357, 358, 359, 360, 361, 362,
363, 364, 365, 694, 366, 367, 107, 368, 369, 370, 371, 372,
373, 374, 375, 376])),
node_pairs_with_labels=((tensor([ 0, 1, 2, ..., 427, 427, 427]), tensor([208, 209, 210, ..., 891, 410, 949])),
tensor([1., 1., 1., ..., 0., 0., 0.])),
node_pairs=(tensor([1970, 1536, 875, 312, 2139, 2280, 837, 306, 211, 9, 1761, 2450,
1650, 531, 1843, 2565, 2585, 2034, 1810, 2371, 742, 606, 125, 1778,
443, 1994, 2168, 70, 2359, 2525, 2263, 2471, 274, 1682, 1120, 1191,
2400, 236, 2396, 365, 1894, 2509, 1798, 1417, 2387, 1115, 379, 2242,
504, 1701, 718, 2512, 2227, 1015, 621, 1533, 1837, 2313, 1153, 1800,
352, 2011, 1810, 2100, 60, 2476, 2238, 1530, 690, 1791, 897, 2202,
90, 447, 2131, 1907, 2125, 1136, 1811, 714, 118, 1804, 1240, 1557,
299, 2222, 1902, 1703, 1035, 1598, 1273, 2354, 189, 2294, 1013, 1503,
2010, 2551, 1703, 105, 306, 1719, 1558, 1758, 2205, 1703, 2249, 910,
1791, 2523, 2451, 1725, 502, 2228, 2612, 401, 244, 1987, 426, 1435,
1478, 1067, 1830, 487, 1848, 2045, 1692, 1154, 1692, 1986, 1228, 565,
2140, 1765, 2238, 88, 1295, 1347, 1954, 1986, 1609, 1810, 979, 2581,
1013, 2332, 1488, 480, 604, 945, 1478, 1952, 1119, 343, 1158, 2242,
597, 171, 1892, 2024, 1592, 1835, 578, 396, 1752, 541, 1135, 232,
1778, 466, 2492, 2030, 2099, 973, 698, 1712, 729, 1419, 1501, 48,
647, 2353, 1725, 1013, 1358, 2036, 1651, 1182, 1416, 480, 1654, 836,
985, 24, 1717, 1158, 2156, 2349, 786, 2516, 659, 1322, 1149, 2024,
842, 1636, 537, 1427, 97, 1529, 848, 2524, 2515, 963, 731, 1867,
1542, 1554, 2180, 2163, 1889, 1004, 743, 525, 1072, 2305, 609, 1388,
735, 341, 2083, 1308, 1903, 267, 517, 928, 703, 1169, 382, 472,
843, 2272, 2115, 626, 321, 1320, 109, 837, 1735, 1696, 1795, 657,
682, 935, 1824, 1755]),
tensor([1702, 2580, 1942, 1610, 24, 279, 2248, 2048, 1909, 723, 223, 1412,
1360, 606, 415, 2567, 1706, 2130, 27, 2367, 1620, 1666, 2363, 350,
1010, 169, 2169, 441, 290, 2220, 94, 2493, 749, 733, 1583, 2339,
963, 153, 2394, 2469, 2034, 2519, 1782, 963, 638, 203, 904, 2539,
2396, 235, 2261, 569, 1630, 1452, 1978, 726, 1330, 873, 1401, 852,
1944, 25, 511, 1184, 1527, 524, 151, 43, 311, 191, 146, 734,
155, 1849, 1463, 210, 1119, 1189, 1810, 148, 2165, 306, 706, 2168,
2386, 284, 448, 1234, 915, 1093, 1111, 577, 1509, 2155, 1849, 2241,
118, 269, 759, 2476, 1805, 1358, 2340, 464, 788, 463, 2248, 139,
1609, 867, 1804, 687, 1976, 2395, 1516, 1096, 565, 1295, 487, 2059,
1902, 59, 889, 2026, 1013, 1448, 1423, 2514, 1472, 1876, 2580, 1398,
1301, 823, 1969, 1309, 1171, 281, 590, 1999, 1912, 1869, 524, 449,
1521, 1481, 1465, 598, 1351, 2388, 694, 2206, 870, 311, 2252, 2241,
1628, 790, 1580, 478, 1351, 2157, 1974, 1217, 1358, 647, 996, 1701,
1624, 1358, 1169, 903, 868, 2153, 1015, 1358, 1221, 529, 1074, 1031,
468, 577, 2596, 1848, 1287, 2034, 306, 559, 1922, 2359, 1567, 2403,
1358, 201, 2560, 589, 324, 2358, 947, 2517, 77, 2265, 1986, 1332,
118, 2139, 1593, 860, 1353, 1908, 228, 1428, 2530, 2399, 2199, 1846,
1176, 2686, 81, 1303, 2130, 1810, 1986, 2057, 342, 1278, 883, 734,
701, 613, 1557, 1627, 2592, 2373, 952, 327, 801, 552, 1203, 1825,
2542, 1441, 2117, 1358, 1450, 2684, 426, 1422, 2324, 1713, 2299, 1740,
1224, 661, 1183, 1731])),
node_features={'feat': tensor([[0.0000, 0.0000, 0.0000, ..., 0.0000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.0000, ..., 0.0000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.0556, ..., 0.0000, 0.0000, 0.0000],
...,
[0.0000, 0.0000, 0.0000, ..., 0.0000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.0000, ..., 0.0000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.0000, ..., 0.0000, 0.0000, 0.0000]])},
negative_srcs=None,
negative_node_pairs=(tensor([[ 0, 0, 0, 0, 0],
[ 1, 1, 1, 1, 1],
[ 2, 2, 2, 2, 2],
...,
[206, 206, 206, 206, 206],
[207, 207, 207, 207, 207],
[427, 427, 427, 427, 427]]),
tensor([[ 377, 378, 379, 380, 381],
[ 382, 383, 876, 384, 385],
[ 386, 719, 387, 388, 389],
...,
[1192, 674, 1306, 456, 1307],
[ 941, 895, 1308, 782, 1309],
[1310, 257, 891, 410, 949]])),
negative_dsts=tensor([[1912, 1606, 731, 560, 71],
[1003, 2558, 2337, 376, 1086],
[1920, 354, 93, 605, 2081],
...,
[1049, 2288, 2448, 126, 1551],
[1638, 1915, 1695, 83, 2333],
[2620, 25, 133, 1661, 1839]]),
labels=None,
input_nodes=tensor([1970, 1536, 875, ..., 1166, 1678, 1364]),
edge_features=[{},
{},
{}],
compacted_node_pairs=(tensor([ 0, 1, 2, 3, 4, 5, 6, 7, 670, 8, 9, 10, 11, 12,
13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 767, 25,
26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53,
54, 55, 56, 57, 58, 59, 17, 60, 61, 62, 63, 64, 65, 66,
67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 334, 77, 78, 79,
80, 721, 81, 658, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91,
658, 92, 7, 93, 94, 95, 96, 658, 97, 98, 66, 99, 100, 101,
730, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114,
115, 116, 115, 333, 117, 118, 119, 120, 63, 121, 122, 845, 123, 333,
124, 17, 653, 125, 88, 752, 126, 127, 128, 129, 109, 130, 131, 132,
133, 45, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145,
22, 146, 147, 148, 867, 149, 150, 151, 813, 152, 153, 154, 155, 156,
101, 88, 694, 157, 158, 159, 778, 127, 160, 161, 162, 163, 164, 133,
165, 166, 789, 861, 167, 168, 169, 137, 170, 171, 172, 173, 174, 175,
176, 177, 689, 178, 379, 179, 180, 181, 182, 183, 184, 801, 185, 186,
187, 188, 189, 190, 191, 192, 826, 193, 194, 195, 196, 643, 788, 197,
198, 199, 940, 946, 200, 201, 202, 203, 952, 6, 795, 204, 205, 454,
911, 206, 207, 427]),
tensor([ 208, 209, 210, 211, 163, 212, 213, 214, 215, 216, 217, 218,
219, 20, 220, 221, 222, 348, 223, 224, 225, 226, 227, 228,
229, 230, 231, 232, 466, 233, 234, 235, 236, 237, 238, 239,
178, 680, 240, 241, 16, 242, 440, 178, 243, 1007, 244, 245,
36, 246, 247, 248, 993, 249, 250, 251, 252, 253, 254, 255,
256, 257, 258, 259, 960, 260, 989, 261, 262, 263, 772, 353,
264, 768, 265, 266, 131, 267, 17, 491, 668, 7, 268, 767,
269, 270, 271, 829, 272, 273, 1040, 1066, 487, 274, 768, 275,
334, 812, 276, 62, 448, 694, 277, 278, 279, 503, 213, 280,
124, 281, 77, 701, 282, 758, 283, 284, 118, 122, 112, 285,
81, 787, 286, 287, 88, 514, 908, 824, 288, 289, 209, 936,
290, 291, 292, 933, 293, 294, 295, 296, 377, 297, 260, 298,
299, 300, 797, 301, 302, 303, 854, 304, 305, 262, 705, 275,
439, 1129, 306, 1097, 302, 307, 308, 309, 694, 155, 310, 47,
311, 694, 197, 400, 312, 313, 51, 694, 314, 315, 316, 317,
318, 1066, 319, 113, 320, 16, 7, 321, 322, 26, 323, 324,
694, 325, 326, 327, 328, 329, 330, 331, 809, 332, 333, 761,
334, 4, 335, 336, 337, 338, 339, 340, 764, 341, 342, 343,
344, 345, 346, 347, 348, 17, 333, 349, 350, 351, 352, 353,
354, 355, 79, 785, 814, 356, 357, 358, 359, 360, 361, 362,
363, 364, 365, 694, 366, 367, 107, 368, 369, 370, 371, 372,
373, 374, 375, 376])),
compacted_negative_srcs=None,
compacted_negative_dsts=tensor([[ 377, 378, 379, 380, 381],
[ 382, 383, 876, 384, 385],
[ 386, 719, 387, 388, 389],
...,
[1192, 674, 1306, 456, 1307],
[ 941, 895, 1308, 782, 1309],
[1310, 257, 891, 410, 949]]),
blocks=[Block(num_src_nodes=2600, num_dst_nodes=2534, num_edges=7690),
Block(num_src_nodes=2534, num_dst_nodes=2282, num_edges=7106),
Block(num_src_nodes=2282, num_dst_nodes=1311, num_edges=3876)],
)
Defining Model for Node Representation¶
Let’s consider training a 2-layer GraphSAGE with neighbor sampling. The model can be written as follows:
[6]:
import dgl.nn as dglnn
import torch.nn as nn
import torch.nn.functional as F
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.hidden_size = hidden_size
self.predictor = nn.Sequential(
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
Defining Traing Loop¶
The following initializes the model and defines the optimizer.
[7]:
in_size = feature.size("node", None, "feat")[0]
model = SAGE(in_size, 128).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
The following is the training loop for link prediction and evaluation.
[8]:
import tqdm
for epoch in range(3):
model.train()
total_loss = 0
for step, data in tqdm.tqdm(enumerate(train_dataloader)):
# Get node pairs with labels for loss calculation.
compacted_pairs, labels = data.node_pairs_with_labels
node_feature = data.node_features["feat"]
# Convert sampled subgraphs to DGL blocks.
blocks = data.blocks
# Get the embeddings of the input nodes.
y = model(blocks, node_feature)
logits = model.predictor(
y[compacted_pairs[0]] * y[compacted_pairs[1]]
).squeeze()
# Compute loss.
loss = F.binary_cross_entropy_with_logits(logits, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
print(f"Epoch {epoch:03d} | Loss {total_loss / (step + 1):.3f}")
38it [00:03, 10.00it/s]
Epoch 000 | Loss 0.562
38it [00:03, 9.61it/s]
Epoch 001 | Loss 0.450
38it [00:03, 9.82it/s]
Epoch 002 | Loss 0.444
Evaluating Performance with Link Prediction¶
[9]:
model.eval()
datapipe = gb.ItemSampler(test_set, batch_size=256, shuffle=False)
# Since we need to use all neghborhoods for evaluation, we set the fanout
# to -1.
datapipe = datapipe.sample_neighbor(graph, [-1, -1])
datapipe = datapipe.fetch_feature(feature, node_feature_keys=["feat"])
datapipe = datapipe.copy_to(device)
eval_dataloader = gb.DataLoader(datapipe, num_workers=0)
logits = []
labels = []
for step, data in tqdm.tqdm(enumerate(eval_dataloader)):
# Get node pairs with labels for loss calculation.
compacted_pairs, label = data.node_pairs_with_labels
# The features of sampled nodes.
x = data.node_features["feat"]
# Forward.
y = model(data.blocks, x)
logit = (
model.predictor(y[compacted_pairs[0]] * y[compacted_pairs[1]])
.squeeze()
.detach()
)
logits.append(logit)
labels.append(label)
logits = torch.cat(logits, dim=0)
labels = torch.cat(labels, dim=0)
# Compute the AUROC score.
from sklearn.metrics import roc_auc_score
auc = roc_auc_score(labels.cpu(), logits.cpu())
print("Link Prediction AUC:", auc)
5it [00:00, 24.08it/s]
Link Prediction AUC: 0.6858300577255676
Conclusion¶
In this tutorial, you have learned how to train a multi-layer GraphSAGE for link prediction with neighbor sampling.