【GNN框架系列】DGL第二讲：使用Deep Graph Library实现GNN进行链接预测

作者：CHEONG

公众号：AI机器学习与知识图谱

研究方向：自然语言处理与知识图谱

本文先简单概述GNN链接预测任务，接下来使用Deep Graph Library实现GNN进行链接预测，并对代码进行详细介绍，若需获取模型的完整代码，可关注公众号后回复：DGL第二讲完整代码

一、GNN链接预测概述

GNN链接预测任务，即预测图中两个节点之间的边是否存在。在Social Recommendation，Knowledge Graph Completion等应用中都需要进行链接预测。模型实现上是将链接预测任务看成一个二分类任务：

1. 将图中存在的边作为正样本；

2. 负采样一些图中不存在的边作为负样本；

3. 将正样例和负样例合并划分为训练集和测试集；

4. 可以采用二分类模型的评估指标来评估模型的效果，例如：AUC值

在一些场景下例如大规模推荐系统或信息检索，模型需要评估top-k预测结果的准确性，因此对于链接预测任务还需要一些其他的评估指标来衡量模型最终效果：

1. MR（MeanRank）

2. MRR（Mean Reciprocal Rank）

3. Hit@n

MR, MRR, Hit@n指标含义：假设整个图谱中共n个实体，评估前先进行如下操作：

（1）将一个正确的三元组$(h,r,t)$中的头实体h或者尾实体t，依次替换成整个图谱中的其他所有实体，这样会产生n个三元组；

（2）对（1）中产生的n个三元组分别计算其能量值，例如在TransE中计算$h+r-t$的值，这样n个三元组分别对应自己的能量值；

（3）对上述n个三元组按照能量值进行升序排序，记录每个三元组排序后的序号；

（4）对所有正确的三元组都进行上述三步操作

MR指标： 将整个图谱中每个正确三元组的能量值排序后的序号取平均得到的值；

MRR指标： 将整个图谱每个正确三元组的能量排序后的序号倒数取平均得到的值；

Hit@n指标： 整个图谱正确三元组的能量排序后序号小于n的三元组所占的比例。

因此对于链接预测任务来说，MR指标越小，模型效果越好；MRR和Hit@n指标越大，模型效果越好。接下来本文将在Cora引文数据集上，预测两篇论文之间是否存在引用关系或被引用关系。

二、GNN链接预测实现

接下来使用DGL框架实现GNN模型进行链接任务，对代码给出详细解释。首先如下所示，先加载需要使用的dgl库和pytorch库；

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

数据加载：下面代码加载dgl库提供的Cora数据对象，dgl库中Dataset数据集可能是包含多个图的，所以加载的dataset对象是一个list，list中的每个元素对应该数据的一个graph，但Cora数据集是由单个图组成，因此直接使用dataset[0]取出graph。

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

正负数据划分：随机抽取数据集中10%的边作为测试集中的正样例，剩下的90%数据集中的边作为训练集，然后随机为训练集和测试集，负采样生成相同数量的负样例，使得训练集和测试集中的正负样本比例为1:1，将数据集中边的集合划分到训练集和测试集中，训练集90%，测试集10%

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())))
adj_neg = 1 - adj.todense() - np.eye(g.number_of_nodes())
neg_u, neg_v = np.where(adj_neg != 0)

neg_eids = np.random.choice(len(neg_u), g.number_of_edges() // 2)
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:]]

注意： 在模型训练时，需要将图中在测试集中10%的边移除掉，防止数据泄露，使用dgl.remove_edges

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

from dgl.nn import SAGEConv

# 定义一个两层的GraphSage模型
class GraphSAGE(nn.Module):
    def __init__(self, in_feats, h_feats):
        super(GraphSAGE, self).__init__()
        self.conv1 = SAGEConv(in_feats, h_feats, 'mean')
        self.conv2 = SAGEConv(h_feats, h_feats, 'mean')

    def forward(self, g, in_feat):
        h = self.conv1(g, in_feat)
        h = F.relu(h)
        h = self.conv2(g, h)
        return h

节点对得分函数定义方式：模型通过定义函数来预测两个节点表示之间的得分，从而来判断两个节点之间存在边的可能性，在GNN节点分类任务重，模型是训练得到单个节点的表征，但在链接计算任务中是预测节点对的表征

注意： 在给出节点对的预测得分函数之前，先需要理解一下：DGL使用方式是先将节点对视为一个图，同时一条边可用来描述一对节点。在链接预测中，会得到一个正图，它包含所有的正例子作为边，以及一个负图，它包含所有的负例子。正图和负图将包含与原始图相同的节点集。这使得在多个图中传递节点特征更容易进行计算。可以直接将在整个图上计算的节点表示形式提供给正图和负图，用于计算节点对的两两得分。

将节点对视为图的好处是可以使用DGLGraph.apply_edges方法，基于节点的特征表示和原始图中边的特征表示可以方便地计算新产生边的特征表示。DGL提供了一组优化的内置函数，可以直接使用原始节点/边特征表示计算新的边特征表示。

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())

官方和自定义两种得分函数实现：接下来定义两个节点之间的得分函数预测，可以直接使用DGL提供的，也可以自定义，下面DotPredictor是官方提供的预测函数，MLPPredictor是自定义的预测函数。

import dgl.function as fn

class DotPredictor(nn.Module):
    def forward(self, g, h):
        with g.local_scope():
            g.ndata['h'] = h
            # 通过源节点特征“h”和目标节点特征“h”之间的点积计算两点之间存在边的Score
            g.apply_edges(fn.u_dot_v('h', 'h', 'score'))
            # u_dot_v为每条边返回一个元素向量，因此需要squeeze操作
            return g.edata['score'][:, 0]

class MLPPredictor(nn.Module):
    def __init__(self, h_feats):
        super().__init__()
        self.W1 = nn.Linear(h_feats * 2, h_feats)
        self.W2 = nn.Linear(h_feats, 1)

    def apply_edges(self, edges):
        """
        Computes a scalar score for each edge of the given graph.
        Parameters
        ----------
        edges :
            Has three members ``src``, ``dst`` and ``data``, each of
            which is a dictionary representing the features of the
            source nodes, the destination nodes, and the edges
            themselves.

        Returns
        -------
        dict
            A dictionary of new edge features.
        """
        h = torch.cat([edges.src['h'], edges.dst['h']], 1)
        return {'score': self.W2(F.relu(self.W1(h))).squeeze(1)}

    def forward(self, g, h):
        with g.local_scope():
            g.ndata['h'] = h
            g.apply_edges(self.apply_edges)
            return g.edata['score']

接下来直接进行模型训练：

optimizer = torch.optim.Adam(itertools.chain(model.parameters(), pred.parameters()), lr=0.01)

all_logits = []
for e in range(100):
    # 前向传播
    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)

    # 后向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if e % 5 == 0:
        print('In epoch {}, loss: {}'.format(e, loss))

# 检测结果准确性
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))

上面是模型的训练函数，和pytorch模型训练过程都是相似的，训练过程如下图所示：

In epoch 0, loss: 0.6172636151313782
In epoch 5, loss: 0.6101921796798706
In epoch 10, loss: 0.5864554047584534
In epoch 15, loss: 0.5405876040458679
In epoch 20, loss: 0.4583510458469391
In epoch 25, loss: 0.39045605063438416
In epoch 30, loss: 0.34702828526496887
In epoch 35, loss: 0.3122958838939667
In epoch 40, loss: 0.2834944725036621
In epoch 45, loss: 0.25488677620887756
In epoch 50, loss: 0.22920763492584229
In epoch 55, loss: 0.20638766884803772
In epoch 60, loss: 0.18289318680763245
In epoch 65, loss: 0.16009262204170227
In epoch 70, loss: 0.1381770521402359
In epoch 75, loss: 0.11725720018148422
In epoch 80, loss: 0.09779688715934753
In epoch 85, loss: 0.07947927713394165
In epoch 90, loss: 0.06309689581394196
In epoch 95, loss: 0.048749890178442
AUC 0.8526520069180836

三、往期精彩

【知识图谱系列】Over-Smoothing 2020综述

【知识图谱系列】基于生成式的知识图谱预训练模型

【知识图谱系列】基于2D卷积的知识图谱嵌入

【知识图谱系列】基于实数或复数空间的知识图谱嵌入

【知识图谱系列】自适应深度和广度图神经网络模型

【知识图谱系列】知识图谱多跳推理之强化学习

【知识图谱系列】知识图谱的神经符号逻辑推理

【知识图谱系列】动态时序知识图谱EvolveGCN

【知识图谱系列】多关系神经网络CompGCN

【机器学习系列】机器学习中的两大学派

干货 | Attention注意力机制超全综述

干货 | NLP中的十个预训练模型

FastText原理和文本分类实战，看这一篇就够了

机器学习算法篇：最大似然估计证明最小二乘法合理性

Word2vec, Fasttext, Glove, Elmo, Bert, Flair训练词向量教程+数据+源码

Word2vec, Fasttext, Glove, Elmo, Bert, Flair训练词向量教程+数据+源码

若需获取模型的完整代码，可关注公众号后回复：DGL第二讲完整代码，有用点个赞呀！