【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