PyTorch图神经网络实践（五）链路预测

链路预测是网络科学里面的一个经典任务，其目的是利用当前已获取的网络数据（包含结构信息和属性信息）来预测网络中会出现哪些新的连边。

本文计划利用networkx包中的网络来进行链路预测，因为目前PyTorch Geometric包中封装的网络还不够多，而很多网络方便用networkx包生成或者处理。

环境配置

首先，安装一个工具包，DeepSNAP。这个包提供了networkx到PyTorch Geometric的接口，可以方便地将networkx中的网络转换成PyTorch Geometric所要求的数据格式。DeepSNAP有两种安装方法：

第一种安装方法

$ pip install deepsnap

第二种安装方法

$ git clone https://github.com/snap-stanford/deepsnap
$ cd deepsnap
$ pip install .

在我服务器上第一种方法报错，使用第二种方法成功了。

其他环境配置参考我之前的系列博文。

链路预测

使用图神经网络进行链路预测包含以下基本步骤：

1. 导入图数据
2. 分割数据集（划分训练边、测试边）
3. 标注正边、采样负边
4. 训练神经网络
5. 测试模型效果

链路预测最开始是一个无监督学习任务，即根据已经看到的网络结构（或者其他属性信息）来推断未知连边是否存在，但是这样的话就比较难以验证。只有在动态网络（或称时序网络）中才会有这样的数据以供实验验证，可以用前一段时间的网络结构来预测后一段时间的网络结构。然而，很多网络没有时间信息，在这样的网络中如何验证呢？

后来，学者提出了用有监督的方式来进行链路预测，也就是将其视为二分类任务，将网络中存在的边都视为正样本（即正边），不存在的连边都当作负样本（即负边）。然后，将这些边分为两部分，一部分为训练集，一部分为测试集。训练集和测试集中都包含正边和负边，目的是在训练集上训练出一个模型能够准确分类这两种边，然后再在测试集上验证效果。

然而，大多数网络都是稀疏的，也就是说存在边的数量差不多是节点数量的几倍左右，而网络中不存在的边的数量差不多是节点数量的平方（在无向网络中，不存在边的数量等于$(n-1)n/2-m$，其中$n$为节点数，$m$为边数)。这样不存边的数量就远远大于存在边的数量，在有监督学习中就意味着负样本远大于正样本，类别极其不平衡。怎么解决这个问题呢？大家很自然地想到了负采样，就是每次训练的时候随机抽取与正样本等比例的负样本，这样就避免了类别不平衡。

训练结束后，就可以用测试集中的正边和负边来验证模型的效果了。

代码解读

完整代码如下

import networkx as nx
from deepsnap.graph import Graph
import torch
import torch.nn.functional as F
from sklearn.metrics import roc_auc_score
from torch_geometric.utils import negative_sampling
from torch_geometric.nn import GCNConv 
from torch_geometric.utils import train_test_split_edges

G = nx.karate_club_graph()
data = Graph(G)
data.num_features = 3
data.x = torch.ones((data.num_nodes, data.num_features), dtype=torch.float32)
data = train_test_split_edges(data)

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = GCNConv(data.num_features, 128)
        self.conv2 = GCNConv(128, 64)

    def encode(self):
        x = self.conv1(data.x, data.train_pos_edge_index)
        x = x.relu()
        return self.conv2(x, data.train_pos_edge_index)

    def decode(self, z, pos_edge_index, neg_edge_index):
        edge_index = torch.cat([pos_edge_index, neg_edge_index], dim=-1)
        logits = (z[edge_index[0]] * z[edge_index[1]]).sum(dim=-1)
        return logits

    def decode_all(self, z):
        prob_adj = z @ z.t()
        return (prob_adj > 0).nonzero(as_tuple=False).t()


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model, data = Net().to(device), data.to(device)
optimizer = torch.optim.Adam(params=model.parameters(), lr=0.01)


def get_link_labels(pos_edge_index, neg_edge_index):
    E = pos_edge_index.size(1) + neg_edge_index.size(1)
    link_labels = torch.zeros(E, dtype=torch.float, device=device)
    link_labels[:pos_edge_index.size(1)] = 1.
    return link_labels


def train():
    model.train()
    neg_edge_index = negative_sampling(
        edge_index=data.train_pos_edge_index, num_nodes=data.num_nodes,
        num_neg_samples=data.train_pos_edge_index.size(1),
        force_undirected=True,
    )
    optimizer.zero_grad()
    z = model.encode()
    link_logits = model.decode(z, data.train_pos_edge_index, neg_edge_index)
    link_labels = get_link_labels(data.train_pos_edge_index, neg_edge_index)
    loss = F.binary_cross_entropy_with_logits(link_logits, link_labels)
    loss.backward()
    optimizer.step()
    return loss


@torch.no_grad()
def test():
    model.eval()
    perfs = []
    for prefix in ["val", "test"]:
        pos_edge_index = data[f'{prefix}_pos_edge_index']
        neg_edge_index = data[f'{prefix}_neg_edge_index']
        z = model.encode()
        link_logits = model.decode(z, pos_edge_index, neg_edge_index)
        link_probs = link_logits.sigmoid()
        link_labels = get_link_labels(pos_edge_index, neg_edge_index)
        perfs.append(roc_auc_score(link_labels.cpu(), link_probs.cpu()))
    return perfs


best_val_perf = test_perf = 0
for epoch in range(1, 11):
    train_loss = train()
    val_perf, tmp_test_perf = test()
    if val_perf > best_val_perf:
        best_val_perf = val_perf
        test_perf = tmp_test_perf
    log = 'Epoch: {:03d}, Loss: {:.4f}, Val: {:.4f}, Test: {:.4f}'
    print(log.format(epoch, train_loss, best_val_perf, test_perf))

z = model.encode()
final_edge_index = model.decode_all(z)

第11行：将networkx中的graph对象转化为torch_geometric的Data对象
第13行：构造节点特征矩阵（原网络不存在节点特征）
第14行：分割训练边集、验证边集（默认占比0.05）以及测试边集（默认占比0.1）
第16-34行：构造一个简单的图卷积神经网络（两层），包含编码（节点嵌入）、解码（分数预测）等操作
第37行：指定设备
第38行：将模型和数据送入设备
第39行：指定优化器
第42-46行：将训练集中的正边标签设置为1，负边标签设置为0
第49-63行：训练函数，每次训练重新采样负边，计算模型损失，反向传播误差，更新模型参数
第67-78行：测试函数，评估模型在验证集和测试集上的预测准确率
第81-89行：训练模型，每次训练完，输出模型在验证集和测试集上的预测准确率
第91-92行：利用训练好的模型计算网络中剩余所有边的分数

上述代码输出为

Epoch: 001, Loss: 0.9718, Val: 0.8889, Test: 0.9388
Epoch: 002, Loss: 0.6701, Val: 0.8889, Test: 0.9388
Epoch: 003, Loss: 0.7293, Val: 0.8889, Test: 0.9388
Epoch: 004, Loss: 0.7009, Val: 0.8889, Test: 0.9388
Epoch: 005, Loss: 0.6880, Val: 0.8889, Test: 0.9388
Epoch: 006, Loss: 0.6905, Val: 0.8889, Test: 0.9388
Epoch: 007, Loss: 0.6906, Val: 0.8889, Test: 0.9388
Epoch: 008, Loss: 0.6835, Val: 0.8889, Test: 0.9388
Epoch: 009, Loss: 0.6779, Val: 0.8889, Test: 0.9388
Epoch: 010, Loss: 0.6788, Val: 0.8889, Test: 0.9388

训练集中的负边是每次随机采样得到的（第51-55行），而验证集和测试集中的负边则在第14行就已经固定了，所以结果中训练集上的loss一直在变化，而验证集和测试集上的得分没有变化。