简介链路预测是复杂网络分析中的关键任务旨在通过已知拓扑发现尚未观测到的潜在关系。面向网络科学、数据挖掘领域的研究学习者这份Python代码包围绕社交网络、信息网络等场景给出了从数据处理到预测评估的完整实现方案。压缩包共23个文件其中10个.py源文件为核心分别承担数据生成与预处理、图结构构建、相似度指标计算、改进RA算法、测试评估与可视化等功能另有项目配置、编译缓存及两张结果示意图整体仅88KB模块划分清晰轻量易读。资源已有2843人学习适合具备Python基础、希望动手实践链路预测算法的读者。代码涵盖局部相似度经典指标与改进的资源分配方法并配有独立评估脚本便于理解精度、召回率等评价指标与交叉验证流程直接替换输入数据或修改算法函数即可在自定义网络上快速验证不同方法的预测效果对复现论文和拓展实验都有参考价值。 链路预测这个方向理论文章一搜一大把各种相似性指标、图嵌入、GNN变体的论文讲得头头是道。但真到自己动手写代码很多朋友就卡住了数据集怎么构造、负样本怎么采、评估指标到底怎么算、为什么我用Node2Vec训出来的嵌入拿去预测边效果还不如一个Common Neighbors。这篇文章我打算直接用代码说话把链路预测从数据准备、baseline实现到GNN模型训练的完整链路走一遍重点讲清楚每一步为什么这么写以及在实测里容易踩的坑。适合谁来读正在做图神经网络相关课题的学生、刚接触图数据挖掘的工程师或者想给自己的推荐系统、风控模型加一个链路预测模块的朋友。本文默认你会Python了解基本的PyTorch用法但不需要你已经精通图神经网络——我会把涉及的核心概念顺带解释清楚。1. 先把任务定义清楚链路预测到底在预测什么链路预测的形式化定义很简单给定一个图 G(V,E)节点集合 V 是固定的边集合 E 是已知的。我们要预测的是那些尚未被观察到、但在未来可能出现的边或者干脆是当前图中不存在的潜在边。放到推荐系统里就是预测用户和商品之间会不会有一条交互边放到蛋白质相互作用网络里就是预测两个蛋白质之间是否存在某种作用关系。但代码实现之前有一个关键问题必须想明白训练数据里的正样本和负样本到底从哪来。1.1 正负样本划分这不是简单的随机切边最常见的做法是把已知边集合 E 划分为训练集和测试集比如随机取 10% 的边作为测试正样本。但有一点很多人会忽略直接从原图里抠掉测试边之后如果你用剩余图去算节点嵌入这个嵌入里已经包含了测试边的邻居结构信息评测的时候会偏乐观后面第4章我会专门说这个问题。负样本的构造更容易踩坑。链路预测里负样本指的是不存在的边但问题在于你没观察到的边不等于它真的不存在可能只是当前数据还没记录到。实操中我们采取一个折中策略——把当前图中不存在的边视为负样本随机采样一部分参与训练和评估就行。下面这段代码展示了最基础的正负样本划分逻辑import numpy as np import networkx as nx from sklearn.model_selection import train_test_split def link_prediction_split(graph, test_frac0.1, neg_ratio1.0, seed42): 将图的边划分为训练集正样本和测试集正样本同时采样负样本。 rng np.random.default_rng(seed) edges np.array(list(graph.edges())) # 划分正样本 pos_train, pos_test train_test_split(edges, test_sizetest_frac, random_stateseed) # 得到所有不存在的边候选集 non_edges np.array(list(nx.non_edges(graph))) # 负样本数量通常与正样本等量或按比例 n_neg_train int(len(pos_train) * neg_ratio) n_neg_test int(len(pos_test) * neg_ratio) neg_idx rng.choice(len(non_edges), sizen_neg_train n_neg_test, replaceFalse) neg_all non_edges[neg_idx] neg_train neg_all[:n_neg_train] neg_test neg_all[n_neg_train:] return pos_train, neg_train, pos_test, neg_test有几个细节值得展开说第一nx.non_edges在边比较稀疏的大图上会非常慢因为它要枚举所有节点对。实际工程里我一般会改成随机采样节点对再过滤掉已有边的方式复杂度低很多。第二neg_ratio设多少合适设成 1.0 意味着负样本和正样本等量。如果设得太大模型会倾向于把所有边都预测为不存在AUC 可能很高但实际没有意义。后面调参部分我会专门讲这个问题。1.2 划分后的图结构处理正样本划分完成后训练用的图应该只用pos_train里的边来构建否则测试边会泄漏到模型训练过程里。train_graph nx.Graph() train_graph.add_nodes_from(graph.nodes()) train_graph.add_edges_from(pos_train)这里还有一个容易被忽略的问题如果原图不是连通图切掉一批边之后某些节点可能变成孤立点。孤立点没有邻居信息基于GNN或Node2Vec的方法几乎无法对它做出有效预测。处理方式通常有两种要么在训练时把这些节点过滤掉要么测试时对孤立点单独用某种fallback策略比如直接用节点属性特征。2. Baseline代码相似性指标为什么能打又为什么不够用链路预测的经典baseline是各种局部相似性指标。别小看它们在不少数据集上这些指标的表现其实相当不错尤其是面对稀疏图的时候。2.1 四种常用指标的代码实现以Common Neighbors共同邻居、Jaccard系数、Adamic-Adar指数AA和Resource AllocationRA为例它们的核心逻辑都是基于两个节点的共同邻居数量区别在于给不同邻居赋予的权重不一样。用NetworkX实现这些指标非常简单import networkx as nx def similarity_scores(graph, node_u, node_v): 计算两个节点之间的多种相似性指标。 common len(list(nx.common_neighbors(graph, node_u, node_v))) jaccard len(list(nx.common_neighbors(graph, node_u, node_v))) / \ len(set(graph.neighbors(node_u)) | set(graph.neighbors(node_v))) # 直接用 networkx 内置方法 aa_score sum(nx.adamic_adar_index(graph, [(node_u, node_v)])) ra_score sum(nx.resource_allocation_index(graph, [(node_u, node_v)])) return { common_neighbors: common, jaccard: jaccard, adamic_adar: aa_score, resource_allocation: ra_score }从代码能看出来这些指标的计算完全基于局部拓扑结构。它们的核心假设是两个节点共享的邻居越多它们之间未来出现边的可能性越大。2.2 在评测集上快速评估baseline评估时要对测试集里的正样本和负样本分别算分然后把两堆分数放到一起算AUC。这里有一个实操细节对于负样本你要算的也是如果它们之间存在边的相似性分数虽然这条边实际不存在。def evaluate_heuristic(graph, pos_test, neg_test, metricadamic_adar): y_true [] y_score [] for u, v in pos_test: y_true.append(1) y_score.append(link_score(graph, u, v, metric)) for u, v in neg_test: y_true.append(0) y_score.append(link_score(graph, u, v, metric)) return roc_auc_score(y_true, y_score)我实测过Cora和Citeseer数据集AA和RA通常比Common Neighbors好一些因为它们在计算共同邻居时对度数高的节点进行了降权避免了大路枢纽型节点带来的虚假高相似度。但这类方法的瓶颈也很明显它们只能利用局部一跳或两跳的信息对网络中的长距离依赖模式无能为力。比如在一个社交网络里两个用户虽然共同好友不多但各自的社交圈子结构非常相似这种情况下局部指标就很难捕捉到了。于是就有了下面要说的图嵌入和GNN方案。3. 图神经网络做链路预测GAE的完整训练代码图神经网络GNN做链路预测的主流框架是GAEGraph Auto-Encoder核心思路分两步Encoder用GCN或GAT把节点编码成低维向量Decoder对目标节点对 (u,v) 计算它们之间存在边的概率常用的decoder就是两个节点向量的内积或者拼接后经过一个MLP这个框架的好处是端到端训练模型能自动从数据里学习什么样的连接模式像是一条真实的边。3.1 用PyTorch Geometric实现GAE如果连PyTorch Geometric都没装先装一下。友情提醒最好用conda安装因为pip直接装源码编译经常会遇到CUDA版本不匹配的问题conda install pyg -c pyg下面是完整的GAE实现代码包含了模型定义、训练循环和评估逻辑。这是我在实际项目里跑通过的版本可以直接拿去用import torch import torch.nn as nn import torch.nn.functional as F from torch_geometric.nn import GCNConv from torch_geometric.utils import train_test_split_edges, negative_sampling class GAEEncoder(nn.Module): 两层GCN编码器把节点编码为低维向量。 def __init__(self, in_dim, hidden_dim, out_dim): super().__init__() self.conv1 GCNConv(in_dim, hidden_dim) self.conv2 GCNConv(hidden_dim, out_dim) def forward(self, x, edge_index): x self.conv1(x, edge_index) x F.relu(x) x F.dropout(x, trainingself.training) x self.conv2(x, edge_index) return x class GAEDecoder(nn.Module): 内积解码器计算节点对之间存在边的概率。 def forward(self, z, edge_index): # 取出每条边两端节点的嵌入 z_u z[edge_index[0]] z_v z[edge_index[1]] # 内积作为相似度分数 return (z_u * z_v).sum(dim-1)训练循环里有两个关键操作边划分和负采样。PyTorch Geometric提供了train_test_split_edges函数它会自动把边划分为train/val/test三部分并且对val和test的负样本预先生成。def train_gae(model, data, optimizer): model.train() optimizer.zero_grad() # 只用训练边的子图来做编码避免信息泄漏 z model(data.x, data.train_pos_edge_index) # 负采样从所有不存在边的节点对中随机采样 neg_edge_index negative_sampling( edge_indexdata.train_pos_edge_index, num_nodesdata.num_nodes, num_neg_samplesdata.train_pos_edge_index.size(1) ) # 正样本分数和负样本分数 pos_score model.decoder(z, data.train_pos_edge_index) neg_score model.decoder(z, neg_edge_index) # 损失正样本尽量接近1负样本尽量接近0 loss F.binary_cross_entropy_with_logits( torch.cat([pos_score, neg_score]), torch.cat([torch.ones_like(pos_score), torch.zeros_like(neg_score)]) ) loss.backward() optimizer.step() return loss.item()3.2 为什么GAE比Node2Vec更适合链路预测很多人一开始用的是Node2Vec或DeepWalk生成的节点嵌入然后同样用内积做链路预测但效果通常不如GAE。原因不复杂Node2Vec这类方法本质上是无监督的它们优化的是让随机游走共现的节点向量更接近这个目标跟链路预测这个任务并不是直接对齐的。而GAE在训练过程中编码器根据当前边的结构信息生成嵌入解码器针对这条边存在与否做优化整个模型是端到端朝着链路预测的目标学的特征自然更贴合任务。但这不意味着Node2Vec没用。在图上节点特征缺失的情况下Node2Vec的嵌入可以作为GNN的初始特征输入两者其实是互补关系。3.3 训练时最常见的两个错误第一个错误用全图的边计算节点嵌入后再划分测试集。这导致的后果是训练和测试高度重叠评估结果虚高看起来AUC 0.95部署到线上立刻打回原形。第二个错误负采样时混入了正样本。也就是说负采样函数里有可能会随机采到本来就存在的边。虽然概率不高但在图比较稠密的时候这个概率不可忽略。一定要确保负样本集合和正样本集合没有交集。negative_sampling函数默认不会采出已有的边但如果你自己实现采样逻辑必须手动检查。4. 评估指标与实验设计AUC、AP和那些容易翻车的细节链路预测的评估指标通常用AUCArea Under ROC Curve和APAverage Precision。AUC衡量的是模型对正负样本排序能力的整体表现AP更关注正样本排在前面有多准。在边预测这种正负样本比例极度不平衡的场景里这两个指标能反映模型的综合能力但前提是评测样本的构造必须合理。4.1 AUC和AP的计算代码from sklearn.metrics import roc_auc_score, average_precision_score def evaluate_link_prediction(model, data, edge_index_pos, edge_index_neg): model.eval() with torch.no_grad(): z model(data.x, data.train_pos_edge_index) pos_score model.decoder(z, edge_index_pos) neg_score model.decoder(z, edge_index_neg) y_true torch.cat([torch.ones_like(pos_score), torch.zeros_like(neg_score)]) y_score torch.cat([pos_score, neg_score]) auc roc_auc_score(y_true.cpu().numpy(), y_score.cpu().numpy()) ap average_precision_score(y_true.cpu().numpy(), y_score.cpu().numpy()) return auc, ap4.2 数据泄漏问题你测试集里的边到底见过没有这个问题我在第3章提过但值得单独拿出来强调。链路预测的评测逻辑和普通 supervised learning 不一样测试边在训练时绝对不能以任何形式出现在编码器看到的图结构里。否则就相当于考试时把答案放在旁边。实操中GAE的训练流程应该是把原始边划分为train_pos_edge_index和val_pos_edge_index、test_pos_edge_index用train_pos_edge_index构建编码器看到的图用这个图生成节点嵌入 z用 z 计算val和test边上的分数PyTorch Geometric的train_test_split_edges已经帮你做好了这套划分但自己做实验时一定要检查数据流别在数据处理环节把测试集边当成已知边传进去。4.3 正负样本不平衡的影响现实中真实图的边数往往远小于不存在的边数。如果把所有不存在的边都作为负样本负样本数量可能是正样本的上百倍。这会带来两个问题模型很容易学到输出很低就对了因为大多数样本确实是负样本评估时AUC可能虚高因为负样本太多随机预测也能得到不错的AUC解决方法是控制负采样比例。我的建议是在[1, 10]范围内调整。也可以用类似困难负样本挖掘的思路训练过程中动态选择那些模型当前最容易判断错的负样本进行优化能明显提升模型对边界情况的判别力。4.4 别忘了看embedding质量链路预测的最终指标当然重要但我习惯在训练完后再做一步检查把节点的embedding用t-SNE或PCA降维到二维空间直接可视化一遍。这一步能直观暴露出很多指标看不出的问题比如训练不充分、特征分布异常、某些类别完全混杂在一起。虽然不是硬性指标但排查问题效率非常高。5. 调参经验与效果兜底方案实测中踩过的几个坑5.1 超参设置的真实体感先说hidden维度。隐藏层维度决定了模型表达能力的上限但也不是越大越好。在Cora这种两千多个节点的小图上hidden_dim设16就能有不错的效果设128反而容易过拟合。我通常的起点是32或64然后观察验证集指标变化趋势。学习率是另一个关键因素。GNN训练里学习率设太大会导致loss震荡不收敛设太小则训练周期很长。Adam优化器配合学习率0.01是比较稳妥的起点。我在几个公开数据集上测试过0.01和0.001两种设置的效果差异能在1-2个点之间浮动。epoch数方面GAE在小图上收敛得很快100-200个epoch基本就稳定了。大图的话可能需要更多。5.2 在Cora上可复现的baseline参考我自己跑的一组实验结果可以给大家做个参照。Cora数据集128维隐藏层负采样比例1:1AUC大约在0.90左右。同样的配置下Adamic-Adar指标大约0.88Node2Vec嵌入加逻辑回归大概0.84。这说明GAE确实能学到一些局部指标捕捉不到的结构模式但优势也没有想象中那么大——这背后反映的是Cora这个数据集的某些特性它的边结构相对规则局部信息已经占了主导。不同数据集上的排名可能有明显变化。在蛋白质相互作用网络上GNN相对启发式指标的提升幅度往往更大因为那些网络里的连接模式更复杂。5.3 不训练模型也能用的兜底方案如果项目时间紧或者图数据量很小直接用启发式指标做baseline也完全可行。我经常建议团队先跑一遍AA或RA指标拿到一个AUC下限后续GNN模型只有超过这个下限才算真的有效。否则模型再花哨也可能只是在拟合噪音。另外提一下SEAL这类基于子图的方法在链路预测上的效果好于GAE因为它不只是看节点本身的嵌入还提取了目标节点对周围的子图结构信息。但它实现起来复杂不少内存开销大。我从实际工程角度建议先用GAE搭通全流程再根据效果决定要不要升级到SEAL。5.4 超参数搜索的简化做法链路预测的超参数空间不算大用网格搜索就能覆盖大部分情况。优先调顺序是负采样比例 隐藏层维度 学习率 drop概率。原因是负采样比例直接影响训练信号的质量这个改动的效果最明显。我自己习惯用optuna做超参搜索设置一个几百次的搜索预算十天就能跑完。每次实验记录AUC、AP和训练时间后面分析起来特别方便。说到给工程落地的建议图数据量大的时候全图GCN的显存开销会非常惊人。这时候可以考虑GraphSAGE这种采样式聚合的模型或者把大规模图切分成子图训练。链路预测这个方向入门的门槛在于理解数据构造的逻辑而不是模型本身——只要正负样本划分和评测流程靠谱模型部分按部就班写就能得到合理的离线指标。我自己每次看到链路预测效果出奇好的报告第一反应都是先去检查它的数据划分是不是出了问题这个习惯帮团队挡掉过好几次毫无意义的模型迭代。本文还有配套的精品资源点击获取