如何用 PyG 的 RandomLinkSplit 为链接预测任务划分边的训练、验证与测试集? 📅 发布时间:2026/9/13 9:49:22 👁 浏览次数: 如何用 PyG 的 RandomLinkSplit 为链接预测任务划分边的训练、验证与测试集【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric做链接预测时监督目标不是节点或整图而是边本身需要把图中的边随机拆成训练、验证、测试三组并且保证测试边不会出现在训练的消息传递边上。PyGPyTorch Geometric用RandomLinkSplit这个 transform 完成这件事输入一个Data或HeteroData对象直接返回(train_data, val_data, test_data)三元组。完整教程见 dataset_splitting.rst可运行的端到端示例见 examples/link_pred.py。RandomLinkSplit 的划分规则按 random_link_split.py 的说明划分遵循一条防泄漏规则训练集不包含验证集和测试集中的边验证集不包含测试集中的边。返回的三个数据对象字段含义如下train_data.edge_index和val_data.edge_index训练和验证阶段用于消息传递的边。训练时只能沿训练边传播信息验证时沿训练边传播test_data.edge_index测试阶段的消息传递边是训练边与验证边的并集*_edge_label_index和*_edge_label默认 key 为edge_label评估和测试用的正、负样本边模型在这批边上做打分。关键参数参数默认值作用num_val0.1验证边数量[0, 1]浮点数表示边比例num_test0.2测试边数量[0, 1]浮点数表示边比例is_undirectedFalse设为True时假设图无向防止反向边的连通性在不同 split 间泄漏只影响图的划分返回的标签数据不会变成无向keyedge_label真实标签属性名。若data[key]不存在会自动创建二分类标签1有边0无边若已存在必须是0到num_classes - 1的类别标签负采样后0表示负边add_negative_train_samplesTrue是否由 transform 添加负训练样本。若模型自己做负采样应设为False否则添加的负样本在各训练迭代间保持不变neg_sampling_ratio1.0负边相对正边数量的采样比例disjoint_train_ratio0.0大于0.0时消息传递与监督所用的训练边互不共享按比例留出作为监督标签edge_types/rev_edge_typesNone操作HeteroData时指定参与划分的边类型及其反向边类型注意is_undirected对二部边类型或edge_type ! rev_edge_types的情况会被忽略。第 1 步用一个小图跑通划分先用 8 个节点、7 条边的图确认 API 行为代码来自教程可直接运行import torch from torch_geometric.data import Data from torch_geometric.transforms import RandomLinkSplit x torch.randn(8, 32) # 节点特征 [num_nodes, num_features] y torch.randint(0, 4, (8, )) # 节点标签 [num_nodes] edge_index torch.tensor([ [2, 3, 3, 4, 5, 6, 7], [0, 0, 1, 1, 2, 3, 4]], ) edge_y torch.tensor([0, 0, 0, 0, 1, 1, 1]) data Data(xx, yy, edge_indexedge_index, edge_yedge_y) edge_transform RandomLinkSplit(num_val0.2, num_test0.2, keyedge_y, is_undirectedFalse, add_negative_train_samplesFalse) train_data, val_data, test_data edge_transform(data)教程中给出的文档示例输出随机划分实际值会不同train_data Data(x[8, 32], edge_index[2, 5], y[8], edge_y[5], edge_y_index[2, 5]) val_data Data(x[8, 32], edge_index[2, 5], y[8], edge_y[2], edge_y_index[2, 2]) test_data Data(x[8, 32], edge_index[2, 6], y[8], edge_y[2], edge_y_index[2, 2])由此可以核对两点val_data的edge_index只含训练边5 条test_data含训练加验证边6 条符合上面的划分规则标签属性名跟随key参数这里传了keyedge_y所以得到edge_y/edge_y_index用默认参数时则是edge_label/edge_label_index。第 2 步应用到真实数据集Planetoid Cora主路径是把RandomLinkSplit放进T.Compose后作为Planetoid的 transform。以下代码取自 examples/link_pred.pyimport os.path as osp import torch import torch_geometric.transforms as T from torch_geometric.datasets import Planetoid if torch.cuda.is_available(): device torch.device(cuda) elif hasattr(torch.backends, mps) and torch.backends.mps.is_available(): device torch.device(mps) else: device torch.device(cpu) transform T.Compose([ T.NormalizeFeatures(), T.ToDevice(device), T.RandomLinkSplit(num_val0.05, num_test0.1, is_undirectedTrue, add_negative_train_samplesFalse), ]) path ./data/Planetoid dataset Planetoid(path, nameCora, transformtransform) # 应用 RandomLinkSplit 后dataset[0] 返回 (train_data, val_data, test_data) train_data, val_data, test_data dataset[0]这里有两个要点is_undirectedTrueCora 这类引用图按无向处理避免反向边跨 split 泄漏连通性。add_negative_train_samplesFalse示例在训练循环里每轮自行负采样所以关闭 transform 内置的负样本添加原因见下节。验证划分结果是否可用示例中用val_data/test_data上的edge_label_index和edge_label计算 AUC 作为评估这本身就是对划分结果的验证——这三个字段齐全说明正负样本已就位import torch from sklearn.metrics import roc_auc_score torch.no_grad() def test(data): model.eval() z model.encode(data.x, data.edge_index) out model.decode(z, data.edge_label_index).view(-1).sigmoid() return roc_auc_score(data.edge_label.cpu().numpy(), out.cpu().numpy())训练时如果模型自行做负采样就在每个 epoch 用torch_geometric.utils.negative_sampling重新生成负边再拼到训练边上做监督示例做法from torch_geometric.utils import negative_sampling neg_edge_index negative_sampling( edge_indextrain_data.edge_index, num_nodestrain_data.num_nodes, num_neg_samplestrain_data.edge_label_index.size(1), methodsparse) 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)) ], dim0)两条负采样路径只能选一种不要叠加add_negative_train_samplesTrue默认时 transform 按neg_sampling_ratio一次性加负样本各训练迭代间不变False时交给训练循环逐轮采样。选哪种取决于模型是否需要每轮新负样本文档给出的判据是模型自己做了负采样就设为False。边界条件与报错对照HeteroData 必须指定edge_types否则抛出ValueError: The RandomLinkSplit transform expects edge_types to be specified when operating on HeteroData objects。边数不够num_val与num_test占掉的边过多、训练边不足或disjoint_train_ratio切完不足时抛出ValueError: Insufficient number of edges for training。负边不足负采样凑不够neg_sampling_ratio要求的数量时会发出警告There are not enough negative edges to satisfy the provided sampling ratio. The ratio will be adjusted to ...并自动按比例调整不会中断。num_val、num_test传整数时按条数计算传[0, 1]浮点数时按比例计算注意两者相加不能超过边数上限否则触发上面的训练边不足错误。延伸节点划分RandomNodeSplit和整图划分的对应做法与本文同处 dataset_splitting.rst。完整的 GCN 链接预测训练与 AUC 评估流程见 examples/link_pred.pyRandomLinkSplit的逐参数实现见 random_link_split.py。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考