基于PyG构建异构图GNN模型实现铁路晚点预测求助

一、模型架构设计（边回归方向） #

你的任务是预测列车到PR节点的边晚点属性，属于数值型边回归任务，核心思路分三步：

• 节点嵌入提取：给列车、PR两种节点类型分别设计特征提取逻辑，利用路网结构（PR→PR边）传递拓扑信息优化PR节点嵌入
• 目标边特征融合：对列车→PR的边，通过拼接/差值/点积等方式融合两端节点的嵌入
• 回归输出：用MLP将融合后的边特征映射为晚点预测值

二、to_hetero适配异构图的正确用法 #

不少人用to_hetero失败，核心是先定义同构图基础模型，再转换为异构模型，同时明确匹配所有节点/边类型：

步骤1：定义同构图基础GNN模块 #

这个模块只处理同构图的节点更新逻辑，后续自动适配异构场景：

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, to_hetero, Linear

class BaseGNN(torch.nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        self.conv1 = GCNConv(-1, hidden_channels)  # -1自动适配输入特征维度
        self.conv2 = GCNConv(hidden_channels, hidden_channels)
        
    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x

步骤2：封装异构边预测模型 #

将基础GNN转为异构版本，再添加针对目标边的预测头：

class HeteroDelayPredictor(torch.nn.Module):
    def __init__(self, hidden_channels, metadata):
        super().__init__()
        # 转换为异构GNN，metadata直接用HeteroData的元数据
        self.gnn = to_hetero(BaseGNN(hidden_channels), metadata=metadata)
        # 边预测头：拼接列车和PR节点的嵌入做回归
        self.predict_head = Linear(2 * hidden_channels, 1)
    
    def forward(self, data):
        # 先更新所有节点的嵌入
        x_dict = self.gnn(data.x_dict, data.edge_index_dict)
        
        # 提取列车→PR边的两端节点嵌入
        edge_index = data.edge_index_dict[('train', 'to', 'pr')]
        train_emb = x_dict['train'][edge_index[0]]
        pr_emb = x_dict['pr'][edge_index[1]]
        
        # 融合后输出预测值
        edge_feat = torch.cat([train_emb, pr_emb], dim=-1)
        return self.predict_head(edge_feat).squeeze()

注：metadata可通过data.metadata()获取，会返回你的HeteroData的节点类型、边类型元组。

三、训练流程实现 #

步骤1：初始化组件 #

# 加载你已构建的数据集
dataset = RailwayGraphDataset(root='./data')
data = dataset[0]

# 初始化模型、优化器、损失函数
model = HeteroDelayPredictor(hidden_channels=128, metadata=data.metadata())
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = torch.nn.MSELoss()  # 回归任务用均方误差

步骤2：训练循环 #

model.train()
for epoch in range(100):
    optimizer.zero_grad()
    
    # 前向传播得到预测值
    pred = model(data)
    # 获取真实标签：假设列车→PR边的晚点属性存在edge_attr_dict中
    target = data.edge_attr_dict[('train', 'to', 'pr')].squeeze()
    
    # 计算损失并反向传播
    loss = criterion(pred, target)
    loss.backward()
    optimizer.step()
    
    if epoch % 10 == 0:
        print(f'Epoch {epoch:03d}, Loss: {loss.item():.4f}')

四、to_hetero失败常见排查点 #

• 基础GNN的forward参数只能是x和edge_index，如果要处理边属性，需在基础模型中加入edge_attr参数，转异构时同步声明
• 检查metadata与HeteroData的节点/边类型完全匹配，比如边类型是('train', 'to', 'pr')而非('train', 'pr')
• 若节点无初始特征，需给每个节点类型初始化随机特征，比如data.x_dict['train'] = torch.randn(num_train_nodes, 16)

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