PyG搭建R-GCN实现链接预测

前言

关于链接预测的介绍以及链接预测中数据集的划分请参考：链接预测中训练集、验证集以及测试集的划分(以PyG的RandomLinkSplit为例)。

1. 数据处理

导入数据：

path = os.path.abspath(os.path.dirname(os.getcwd())) + '\data\DBLP'
dataset = DBLP(path)
graph = dataset[0]
print(graph)

输出如下：

HeteroData(
  author={
    x=[4057, 334],
    y=[4057],
    train_mask=[4057],
    val_mask=[4057],
    test_mask=[4057]
  },
  paper={ x=[14328, 4231] },
  term={ x=[7723, 50] },
  conference={ num_nodes=20 },
  (author, to, paper)={ edge_index=[2, 19645] },
  (paper, to, author)={ edge_index=[2, 19645] },
  (paper, to, term)={ edge_index=[2, 85810] },
  (paper, to, conference)={ edge_index=[2, 14328] },
  (term, to, paper)={ edge_index=[2, 85810] },
  (conference, to, paper)={ edge_index=[2, 14328] }
)

可以发现，DBLP数据集中有作者(author)、论文(paper)、术语(term)以及会议(conference)四种类型的节点。DBLP中包含14328篇论文(paper)， 4057位作者(author)， 20个会议(conference)， 7723个术语(term)。作者分为四个领域：数据库、数据挖掘、机器学习、信息检索。

由于conference节点没有特征，因此需要预先设置特征：

graph['conference'].x = torch.ones((graph['conference'].num_nodes, 1))

所有conference节点的特征都初始化为[1]。

利用PyG封装的RandomLinkSplit我们很容易实现数据集的划分：

train_data, val_data, test_data = T.RandomLinkSplit(
        num_val=0.1,
        num_test=0.1,
        is_undirected=True,
        add_negative_train_samples=False,
        disjoint_train_ratio=0,
        edge_types=[('author', 'to', 'paper'), ('paper', 'to', 'term'),
                    ('paper', 'to', 'conference')],
        rev_edge_types=[('paper', 'to', 'author'), ('term', 'to', 'paper'),
                        ('conference', 'to', 'paper')]
    )(graph.to_homogeneous())

最终我们得到train_data, val_data, test_data。

输出一下原始数据集和三个被划分出来的数据集：

Data(node_type=[26128], edge_index=[2, 239566], edge_type=[239566])
Data(node_type=[26128], edge_index=[2, 191654], edge_type=[191654], edge_label=[95827], edge_label_index=[2, 95827])
Data(node_type=[26128], edge_index=[2, 191654], edge_type=[191654], edge_label=[23956], edge_label_index=[2, 23956])
Data(node_type=[26128], edge_index=[2, 215610], edge_type=[215610], edge_label=[23956], edge_label_index=[2, 23956])

从上到下依次为原始数据集、训练集、验证集以及测试集。其中，训练集中一共有95827个正样本，验证集和测试集中均为11978个正样本+11978个负样本。

2. R-GCN链接预测

本次实验使用R-GCN来进行链接预测：首先利用R-GCN对训练集中的节点进行编码，得到节点的向量表示，然后使用这些向量表示对训练集中的正负样本（在每一轮训练时重新采样负样本）进行有监督学习，具体来讲就是利用节点向量求得样本中节点对的内积，然后与标签求损失，最后反向传播更新参数。

2.1 负采样

链接预测训练过程中的每一轮我们都需要对训练集进行采样以得到与正样本数量相同的负样本，验证集和测试集在数据集划分阶段已经进行了负采样，因此不必再进行采样。

负采样函数：

def negative_sample():
    # 从训练集中采样与正边相同数量的负边
    neg_edge_index = negative_sampling(
        edge_index=train_data.edge_index, num_nodes=train_data.num_nodes,
        num_neg_samples=train_data.edge_label_index.size(1), method='sparse')
    # print(neg_edge_index.size(1))   # 3642条负边，即每次采样与训练集中正边数量一致的负边
    edge_label_index = torch.cat(
        [train_data.edge_label_index, neg_edge_index],
        dim=-1,
    )
    edge_label = torch.cat([
        train_data.edge_label,
        train_data.edge_label.new_zeros(neg_edge_index.size(1))
    ], dim=0)

    return edge_label, edge_label_index

这里用到了negative_sampling方法，其参数有：

具体来讲，negative_sampling方法利用传入的edge_index参数进行负采样，即采样num_neg_samples条edge_index中不存在的边。num_nodes指定节点个数，method指定采样方法，有sparse和dense两种方法。

采样后将neg_edge_index与训练集中原有的正样本train.edge_label_index进行拼接以得到完整的样本集，同时也需要在原本的train_data.edge_label后面添加指定数量的0用于表示负样本。

2.2 模型搭建

R-GCN链接预测模型搭建如下：

class RGCN_LP(nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super(RGCN_LP, self).__init__()
        self.conv1 = RGCNConv(in_channels, hidden_channels,
                              num_relations=num_relations, num_bases=30)
        self.conv2 = RGCNConv(hidden_channels, out_channels,
                              num_relations=num_relations, num_bases=30)
        self.lins = torch.nn.ModuleList()
        for i in range(len(node_types)):
            lin = nn.Linear(init_sizes[i], in_channels)
            self.lins.append(lin)

    def trans_dimensions(self, g):
        data = copy.deepcopy(g)
        for node_type, lin in zip(node_types, self.lins):
            data[node_type].x = lin(data[node_type].x)

        return data

    def encode(self, data):
        data = self.trans_dimensions(data)
        homogeneous_data = data.to_homogeneous()
        # print(homogeneous_data)
        edge_index, edge_type = homogeneous_data.edge_index, homogeneous_data.edge_type
        x = self.conv1(homogeneous_data.x, edge_index, edge_type)
        x = self.conv2(x, edge_index, edge_type)
        return x

    def decode(self, z, edge_label_index):
        # z所有节点的表示向量
        src = z[edge_label_index[0]]
        dst = z[edge_label_index[1]]
        # print(dst.size())
        r = (src * dst).sum(dim=-1)
        # print(r.size())
        return r

    def forward(self, data, edge_label_index):
        z = self.encode(data)
        return self.decode(z, edge_label_index)

编码器由一个两层R-GCN组成，用于得到训练集中节点的向量表示，解码器用于得到训练集中节点对向量间的内积。

这里需要注意的是，DBLP转为同质图后虽然有了edge_index和edge_type，但没有所有节点的特征x，这是因为在将异质图转为同质图的过程中，只有所有节点的特征维度相同时才能将所有节点的特征进行合并。因此，我们首先需要将所有节点的特征转换到同一维度(这里以128为例)：

def trans_dimensions(self, g):
    data = copy.deepcopy(g)
    for node_type, lin in zip(node_types, self.lins):
        data[node_type].x = lin(data[node_type].x)

    return data

转换后的data中所有类型节点的特征维度都为128，然后再将其转为同质图：

data = self.trans_dimensions(data)
homogeneous_data = data.to_homogeneous()
Data(node_type=[26128], x=[26128, 128], edge_index=[2, 239566], edge_type=[239566])

此时，我们就可以将homogeneous_data输入到RGCNConv中：

x = self.conv1(homogeneous_data.x, edge_index, edge_type)
x = self.conv2(x, edge_index, edge_type)

由前面可知训练集中的正样本数量为95827，经过负采样函数negative_sample得到95827个负样本，一共191654个样本，最终解码器返回191654个节点对间的内积。

2.3 模型训练/测试

参考前面：PyG搭建GCN实现链接预测

训练300轮，最终测试集上的AUC为0.91：

final best auc: 0.9142443020642201