PyTorch Geometric RandomLinkSplit拆分异常及索引越界问题求助

一、修复数据集拆分不同步问题 #

你的核心问题是RandomLinkSplit参数设置有误，导致edge_index、edge_attr未与edge_label同步拆分。调整参数并遵循正确用法即可解决：

关键参数修正 #

• 指定edge_label_key：明确告知拆分器要基于哪个字段拆分标签，确保edge_attr和edge_index同步关联拆分后的标签边。
• 调整split_labels为True：当需要将标签边拆分为训练/验证/测试集时，该参数需设为True（默认值），否则拆分器只会处理edge_label而忽略关联的边属性和索引。
• 拆分前移回CPU处理：PyTorch Geometric的部分变换在GPU上可能出现同步问题，建议先将Data对象移到CPU拆分，再放回设备。

修正后的拆分代码 #

# 将数据移回CPU进行拆分（避免GPU同步问题）
pyGData = pyGData.cpu()

# 初始化正确的拆分器
split = transforms.RandomLinkSplit(
    is_undirected=True,
    split_labels=True,  # 改为True，拆分标签边
    num_val=0.2,
    num_test=0.2,
    edge_label_key="edge_label"  # 指定标签字段
)

# 执行拆分
train_data, val_data, test_data = split(pyGData)

# 拆分完成后再移回设备
train_data = train_data.to(device)
val_data = val_data.to(device)
test_data = test_data.to(device)

拆分后的数据结构说明 #

拆分后每个子集的结构会包含：

• edge_index：训练/验证/测试用的边索引（对应保留的边）
• edge_attr：对应边的属性
• edge_label：对应边的标签
• edge_label_index：拆分器自动生成的标签边索引（用于明确待预测的边）

二、解决edge索引越界问题 #

索引越界通常由两个原因导致，逐一排查修复：

1. 节点ID非连续从0开始 #

PyTorch Geometric要求节点ID必须是0到num_nodes-1的连续整数，如果原始数据中节点ID是离散业务ID，会直接导致索引越界。

修复方法：重新映射节点ID为连续索引

# 提取所有唯一节点ID
nodes_list = torch_edges.flatten().unique().tolist()
# 创建ID映射字典：原始ID -> 连续索引
node_id_map = {old_id: new_id for new_id, old_id in enumerate(nodes_list)}
# 重新映射edge_index中的节点ID
torch_edges = torch.tensor([[node_id_map[id] for id in row] for row in torch_edges.tolist()], dtype=torch.long).t().contiguous()
# 更新num_nodes为映射后的节点总数
num_nodes = len(nodes_list)

2. num_nodes设置错误 #

确保num_nodes等于实际存在的节点数量，而非原始数据的行数或其他统计值。使用上述映射后的num_nodes即可避免该问题。

3. 无向边重复处理 #

如果原始数据中已经包含双向边（比如A→B和B→A），设置is_undirected=True会自动去重，避免重复边导致的索引混乱；如果是单向边，拆分器会自动生成反向边，无需手动处理。

最终整合代码示例 #

import torch
from torch_geometric.data import Data
from torch_geometric import transforms

# --- 数据预处理修正 ---
edges_np = edges.collect()
torch_edges = torch.tensor(edges_np, dtype=torch.long).t().contiguous()

# 映射节点ID为连续索引
nodes_unique = torch_edges.flatten().unique()
node_id_map = {old: new for new, old in enumerate(nodes_unique)}
torch_edges = torch.tensor([[node_id_map[src], node_id_map[dst]] for src, dst in edges_np], dtype=torch.long).t().contiguous()
num_nodes = len(nodes_unique)

# 转换标签和属性
torch_edges_labels = torch.tensor(df.select("Is Laundering").collect(), dtype=torch.long)
torch_edges_features = torch.tensor(df_features.collect(), dtype=torch.float)

# 创建Data对象
pyGData = Data(
    num_nodes=num_nodes,
    edge_index=torch_edges,
    edge_attr=torch_edges_features,
    edge_label=torch_edges_labels
)

# --- 数据集拆分修正 ---
split = transforms.RandomLinkSplit(
    is_undirected=True,
    split_labels=True,
    num_val=0.2,
    num_test=0.2,
    edge_label_key="edge_label"
)

train_data, val_data, test_data = split(pyGData)

# 移到设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_data = train_data.to(device)
val_data = val_data.to(device)
test_data = test_data.to(device)

内容的提问来源于stack exchange，提问作者Mahana