在之前的介紹中，我們已經學習了使用GNN進行節點分類，比如預測一個圖中的節點所屬的類別。這一節中我們將教你如何進行鏈路預測，比如預測任意兩個節點之間是不是存在邊。

本節你將學到：

• 構建一個GNN的鏈路預測模型
• 在一個小的DGL數據集上訓練和評估模型

鏈路預測

在很多應用中，例如社交推薦、商品推薦以及知識圖譜補全中都存在鏈路預測，就是判斷兩個節點之間是不是存在一條邊。本節將使用論文引用關系數據集，判斷兩篇論文是否存在引用關系。

這個教程將鏈路預測定義為一個二分類的問題：

• 將圖中的每個邊都視為正樣本
• 對不存在邊的節點進行負采樣
• 將正樣本和負樣本都放入訓練集和測試集
• 用任意一個二分類器對模型進行評價，然后計算Area Under Curve（AUC）

在一些領域尤其是大規模推薦系統或信息檢索系統中，你可以能更喜歡Top-k性能指標。

加載graph和features

import dgl.datadataset = dgl.data.CoraGraphDataset() g = dataset[0]

輸出：

NumNodes: 2708NumEdges: 10556NumFeats: 1433NumClasses: 7NumTrainingSamples: 140NumValidationSamples: 500NumTestSamples: 1000 Done loading data from cached files.

設置訓練集和測試集

我們使用10%的邊作為測試集的正樣本，剩下的作為訓練樣本。同時在訓練集和測試集眾采樣相同數量的負樣本。

import dgl import torch import torch.nn as nn import torch.nn.functional as F import itertools import numpy as np import scipy.sparse as spdataset = dgl.data.CoraGraphDataset() g = dataset[0]u, v = g.edges() eids = np.arange(g.number_of_edges()) eids = np.random.permutation(eids) # 將順序打亂test_size = int(len(eids) * 0.1) train_size = g.number_of_edges() - test_size test_pos_u, test_pos_v = u[eids[:test_size]], v[eids[:test_size]] train_pos_u, train_pos_v = u[eids[test_size:]], v[eids[test_size:]]adj = sp.coo_matrix((np.ones(len(u)), (u.numpy(), v.numpy()))) # 利用一個全1向量，以及對應的u和v來構造鄰接矩陣 adj_neg = 1 - adj.todense() - np.eye(g.number_of_nodes()) # 獲得負采樣鄰接矩陣 adj.todense()表示將稀疏矩陣adj變為稠密矩陣 neg_u, neg_v = np.where(adj_neg != 0) # 在鄰接矩陣上進行負采樣neg_eids = np.random.choice(len(neg_u), g.number_of_edges()) test_neg_u, test_neg_v = neg_u[neg_eids[:test_size]], neg_v[neg_eids[:test_size]] train_neg_u, train_neg_v = neg_u[neg_eids[test_size:]], neg_v[neg_eids[test_size:]]

其實上面這段負采樣的方法有一點問題，那就是在全局圖上進行了負采樣，這樣就將測試集中的數據也進行了負采樣，理論上講應該在訓練集上進行負采樣，對測試集不可見。

在訓練過程中，需要從graph中刪除測試集的邊，使用dgl.remove_edges。
dgl.remove_edges通過從父圖中創建一個子圖，這會產生一個復制的過程，因此在大圖中會很慢。因此最好將train graph和test graph保存到本地硬盤方便預處理。

train_g = dgl.remove_edges(g, eids[:test_size])

定義一個GraphSAGE模型

這里構建一個兩層GraphSAGE模型，每一層聚合節點周圍的鄰居信息。DGL提供dgl.nn.SAGEConv創建一個Layer。

from dgl.nn import SAGEConv# ----------- 2. create model -------------- # # build a two-layer GraphSAGE model class GraphSAGE(nn.Module):def __init__(self, in_feats, h_feats):super(GraphSAGE, self).__init__()self.conv1 = SAGEConv(in_feats, h_feats)self.conv2 = SAGEConv(h_feats, h_feats)def forward(self, g, in_feat):h = self.conv1(g, in_feat)h = F.relu(h)h = self.conv2(g, h)return h

模型將會通過一個MLP或者點乘方法來計算兩個節點是否存在邊的概率值：

正采樣graph，負采樣graph

DGL建議我們將一系列節點對the pairs of nodes視為一個graph，因為我們是根據邊來構建的節點對。在鏈路預測中，你擁有一個正采樣節點對所對應的graph，同時也有負采樣節點對所對應的graph。正采樣的graph和負采樣的graph中的節點都包含所有的原始圖中的節點，這會對聚合周圍鄰居信息提供遍歷。你可以直接將節點都喂給正采樣graph和負采樣graph來計算pair-wise scores。

下面的代碼分別用訓練集與測試集構建了一個正采樣graph和一個負采樣graph。

train_pos_g = dgl.graph((train_pos_u, train_pos_v), num_nodes=g.number_of_nodes()) train_neg_g = dgl.graph((train_neg_u, train_neg_v), num_nodes=g.number_of_nodes())test_pos_g = dgl.graph((test_pos_u, test_pos_v), num_nodes=g.number_of_nodes()) test_neg_g = dgl.graph((test_neg_u, test_neg_v), num_nodes=g.number_of_nodes())

將節點對視為一個graph就可以利用DGL.apply_edges方法方便的計算節點的特征了。DGL提供了一組優化方法可以根據原始節點或邊的特征計算新的邊特征。例如dgl.function.u_dot_v計算每一條邊對應節點對的點積。

import dgl.function as fnclass DotPredictor(nn.Module):def forward(self, g, h):with g.local_scope():g.ndata['h'] = h# Compute a new edge feature named 'score' by a dot-product between the# source node feature 'h' and destination node feature 'h'.g.apply_edges(fn.u_dot_v('h', 'h', 'score'))# u_dot_v returns a 1-element vector for each edge so you need to squeeze it.return g.edata['score'][:, 0]

如果業務邏輯比較復雜可以進行重寫。例如下面的模型對每個節點對的特征拼接在一起然后傳遞給MLP。

訓練

在定義了節點表示以及邊的計算方法之后，就可以進行損失函數的計算，以及評價模型了。

損失函數使用的一個簡單的二分交叉熵損失：

評價方法使用AUC：

model = GraphSAGE(g.ndata['feat'].shape[1], 16, 7) # You can replace DotPredictor with MLPPredictor. #pred = MLPPredictor(16) pred = DotPredictor()def compute_loss(pos_score, neg_score):scores = torch.cat([pos_score, neg_score])labels = torch.cat([torch.ones(pos_score.shape[0]), torch.zeros(neg_score.shape[0])])return F.binary_cross_entropy_with_logits(scores, labels)def compute_auc(pos_score, neg_score):scores = torch.cat([pos_score, neg_score]).numpy()labels = torch.cat([torch.ones(pos_score.shape[0]), torch.zeros(neg_score.shape[0])]).numpy()return roc_auc_score(labels, scores)# ----------- 3. set up loss and optimizer -------------- # # in this case, loss will in training loop optimizer = torch.optim.Adam(itertools.chain(model.parameters(), pred.parameters()), lr=0.01)# ----------- 4. training -------------------------------- # all_logits = [] for e in range(100):# forwardh = model(train_g, train_g.ndata['feat'])pos_score = pred(train_pos_g, h)neg_score = pred(train_neg_g, h)loss = compute_loss(pos_score, neg_score)# backwardoptimizer.zero_grad()loss.backward()optimizer.step()if e % 5 == 0:print('In epoch {}, loss: {}'.format(e, loss))# ----------- 5. check results ------------------------ # from sklearn.metrics import roc_auc_score with torch.no_grad():pos_score = pred(test_pos_g, h)neg_score = pred(test_neg_g, h)print('AUC', compute_auc(pos_score, neg_score))

總結

以上是生活随笔為你收集整理的DGL教程【四】使用GNN进行链路预测的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。