DGL教程【四】使用GNN进行链路预测

在之前的介绍中，我们已经学习了使用GNN进行节点分类，比如预测一个图中的节点所属的类别。这一节中我们将教你如何进行链路预测，比如预测任意两个节点之间是不是存在边。

本节你将学到：

• 构建一个GNN的链路预测模型
• 在一个小的DGL数据集上训练和评估模型

链路预测

在很多应用中，例如社交推荐、商品推荐以及知识图谱补全中都存在链路预测，就是判断两个节点之间是不是存在一条边。本节将使用论文引用关系数据集，判断两篇论文是否存在引用关系。

这个教程将链路预测定义为一个二分类的问题：

• 将图中的每个边都视为正样本
• 对不存在边的节点进行负采样
• 将正样本和负样本都放入训练集和测试集
• 用任意一个二分类器对模型进行评价，然后计算Area Under Curve（AUC）

在一些领域尤其是大规模推荐系统或信息检索系统中，你可以能更喜欢Top-k性能指标。

加载graph和features

import dgl.data

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

输出：

  NumNodes: 2708
  NumEdges: 10556
  NumFeats: 1433
  NumClasses: 7
  NumTrainingSamples: 140
  NumValidationSamples: 500
  NumTestSamples: 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 sp

dataset = 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 fn

class 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):
    # forward
    h = 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)

    # backward
    optimizer.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))