GNN图神经网络代码实战：PyG消息传递与节点分类全解析

简介一份覆盖图神经网络核心模块的完整GNN代码包面向想从零实现节点分类、链接预测等任务的机器学习学习者与研究人员。代码包含网络架构定义、数据预处理、训练与评估流程并重点展示了节点嵌入embedding和边edge信息的处理方式便于理解GNN如何聚合邻居信息并生成节点表示。包内共384个文件主体为322个json配置与数据文件可支撑图数据加载和参数设置另有59个zbak备份文件、1个txt说明文件及一个内置zip子包整体压缩包仅2.16MB体量轻、结构清晰。目前已有238人学习或下载适合用来对照论文搭建原型、修改实验或作为课设参考。借助这套代码可以快速上手图数据管道搭建省去从零调试的麻烦重点关注图卷积等关键层的实现细节。1. “gnn图神经网络代码完整”到底指什么一次跑通的消息传递范式想从零复现一个图神经网络gnn图神经网络代码完整这件事难点从来不是“会用 torch”而是“图和神经网络之间那层关系怎么搭起来”。在推荐、风控、分子性质预测这类任务里数据天然是图结构用户和商品的交互是边化合物里原子之间的化学键是边社交关系更不必说。传统机器学习把样本当独立个体等于把这张网拍扁了再建模丢掉了“邻居”这个最值钱的信息。GNN 做的事情反过来——每个节点的表示由它自己的特征和邻居的特征共同决定这种机制叫消息传递。这篇文章面向的是已经会用 PyTorch 写普通 MLP、但第一次碰图数据的工程师。我会先带你拆一个最小可运行的完整代码骨架再看 GCN、GraphSAGE、GAT 三种模型各自的选型逻辑最后给出训练、评估和排障的具体参数。按这套路径走你不需要去啃厚厚的基础设施源码半天之内就能在自己数据集上跑出第一版模型。2. 拆开一份可跑的 GNN 代码数据、邻接矩阵与消息传播主线2.1 用 PyG 组织图数据edge_index 是绕不开的坐标目前从业者搭 GNN 最常见的是两条路PyTorch GeometricPyG和 Deep Graph LibraryDGL。两者都能完成消息传递但风格差异很大。PyG 的 API 更贴近 PyTorch 原生习惯Data 对象就是一个透明的容器DGL 则更像一个独立的图计算引擎。如果只是做实验验证想法我一般先用 PyG代码量更小调试路径更直接。PyG 里一张图就是一个Data对象核心字段只有三个x是节点特征矩阵形状为[N, F]edge_index是边列表形状为[2, E]第一行是起点第二行是终点y是标签。注意edge_index不是邻接矩阵本身它更像边的坐标记录——这个设计决定了后面所有消息传递的效率。import torch from torch_geometric.data import Data # 3 个节点每个节点带 4 维特征 x torch.randn(3, 4) # 两条有向边0-11-2 # 在无向图里每次添加一条边必须同时写 1-0否则消息只单向传播 edge_index torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtypetorch.long) y torch.tensor([0, 1, 0]) data Data(xx, edge_indexedge_index, yy) print(data)这里有个新手最容易犯错的地方edge_index被设计为只存有向边而真实场景里 PyG 内置的 Cora、CiteSeer 这些数据集给你的是无向图。虽然数据集内部通常已经帮你把两条方向的边都存进去了但如果你从 CSV 里自己构造图只往edge_index里塞了(u, v)却没有塞(v, u)训练出来的模型收敛会非常奇怪——因为信息只能从 v 流向 u图变成了严格的单向传递。自己构造数据时用torch_geometric.utils.to_undirected(edge_index)可以直接把边补全。2.2 第一个 GCN 模型卷积在哪里发生GCN 的数学表达并不复杂每一层对节点的邻居特征做一次对称归一化聚合再经过线性变换和激活。用 PyG 的GCNConv写出来代码结构比理论公式直白得多。import torch.nn as nn import torch.nn.functional as F from torch_geometric.nn import GCNConv class GCN(nn.Module): def __init__(self, in_channels, hidden_channels, out_channels): super().__init__() # 两层图卷积聚合邻居特征并做线性变换 self.conv1 GCNConv(in_channels, hidden_channels) self.conv2 GCNConv(hidden_channels, out_channels) def forward(self, x, edge_index): # 第一层卷积 ReLU Dropout x self.conv1(x, edge_index) x F.relu(x) x F.dropout(x, p0.5, trainingself.training) # 第二层输出 logits不接激活 x self.conv2(x, edge_index) return xGCNConv内部默认做了三件事把edge_index转成稀疏邻接矩阵加上自环add_self_loopsTrue然后做对称归一化。归一化的公式是 D^{-1/2} A D^{-1/2}也就是把节点的度考虑进去防止高度数节点聚合出来的数值太小、低度数节点数值太大。理解不了这个归一化没关系但你必须知道存在这一步——它直接决定了为什么 GCN 不需要像图像 ConvNet 那样做复杂的输入标准化。dropout只在训练时生效self.training这个参数由 PyTorch 在调用model.train()/model.eval()时自动切换。这里要注意的是GNN 的 dropout 位置有讲究只在层间加不要在输入特征上强烈丢弃——图数据里特征维度通常已经不大丢掉太多等于把节点信息直接抹掉。2.3 在 Cora 上跑通第一个训练脚本Cora 是 GNN 领域最常用的论文引用数据集2708 个节点、5409 条无向边、每个节点是一个词包向量。PyG 内置了它不需要手动下载数据文件。from torch_geometric.datasets import Planetoid # 第一次运行会自动下载后续会缓存到本地 dataset Planetoid(root./data/Cora, nameCora) data dataset[0] # 训练、验证、测试的节点索引由数据集给定 train_mask data.train_mask val_mask data.val_mask test_mask data.test_mask print(f节点数: {data.num_nodes}, 边数: {data.num_edges}) print(f特征维度: {dataset.num_features}, 类别数: {dataset.num_classes}) print(f训练集: {train_mask.sum().item()}, f验证集: {val_mask.sum().item()}, f测试集: {test_mask.sum().item()})PyG 的Planetoid数据集里已经标记好了需要忽略的自环边并且划分了固定的 train/val/test mask这是它和普通机器学习数据集最大的不同点不是随机抽样而是固定节点索引。因为图上的样本不是独立同分布的如果随意随机划分节点与其邻居的标签会同时出现在训练和测试集里这个信息泄漏会让你的测试结果虚高。很多人在自己数据上复现不出论文指标一半原因是自己用了random_split。训练循环和普通 PyTorch 几乎一样import torch from torch_geometric.datasets import Planetoid device torch.device(cuda if torch.cuda.is_available() else cpu) dataset Planetoid(root./data/Cora, nameCora) data dataset[0].to(device) model GCN(dataset.num_features, 16, dataset.num_classes).to(device) optimizer torch.optim.Adam(model.parameters(), lr0.01, weight_decay5e-4) criterion torch.nn.CrossEntropyLoss() model.train() for epoch in range(200): optimizer.zero_grad() out model(data.x, data.edge_index) # 只计算有标签且被 mask 选中的节点 loss criterion(out[data.train_mask], data.y[data.train_mask]) loss.backward() optimizer.step() if epoch % 10 0: print(fEpoch {epoch}, Loss: {loss.item():.4f})out的 shape 是[2708, 7]每个节点一行分类 logits。训练时只用train_mask选中的 140 个节点的损失去反传这是 GNN 半监督学习的核心设定——用少量标注数据在整张图上传播信息。所以损失计算的边界必须写对漏掉 mask 等于把所有无标签节点也当成了训练样本模型不崩才怪。2.4 评估准确率、F1、还有全图推理训练结束后在测试 mask 上计算准确率。这里推荐直接算 macro F1 而不只是准确率——GNN 常用于类别不均衡的场景准确率会被多数类带偏。from sklearn.metrics import accuracy_score, f1_score model.eval() with torch.no_grad(): out model(data.x, data.edge_index).cpu() pred out.argmax(dim1) test_pred pred[data.test_mask] test_true data.y[data.test_mask] acc accuracy_score(test_true, test_pred) macro_f1 f1_score(test_true, test_pred, averagemacro) print(fTest Accuracy: {acc:.4f}, Macro F1: {macro_f1:.4f})Cora 上两层 GCN、隐藏层 16 维、200 epoch、Adam lr0.01能复现出 0.81 左右的准确率。这组超参是 GCN 原始论文里的设定历代研究者反复验证过可以作为你验证环境是否装对的“基准测试”。如果跑出来只有 0.6 甚至更低先不要怀疑模型回去检查train_mask和数据集的加载是否完整。3. 模型选型GCN、GraphSAGE、GAT 在什么场景下替换3.1 GCN 的边界过度平滑与固定权重GCN 的核心问题在于它只能用同一种权重去对待所有邻居而且层数一深就会出现“过度平滑”。我在一个知识图谱的实体分类项目里试过 5 层 GCN训练集损失照样下降但验证集准确率从第 2 层开始就一路走低。原因是每个节点经过多次聚合后特征越来越接近全局平均失去了身份辨识度。这是 GCN 的结构性缺陷和调参无关。GCN 适合的边界很明确图规模不大、特征维度不高、节点标签和结构的相关性较强、不需要显式建模邻居重要性。常见的做法是在需要为邻居分配不同权重时切换到 GAT在需要采样大图时切换到 GraphSAGE。如果你不确定自己该用哪个先跑 GCN 基线——它仍然是大多数 GNN 项目的起点和最稳的对照模型。在一个用户关系链的欺诈检测场景里我遇到过另一个 GCN 的实际限制边的可靠性不一样。注册设备相同的账号和共用 IP 的账号虽然都是“有边”但后者欺诈置信度明显更高。GCN 不会区分这两条边只能靠模型自己从特征里学出来。这就是为什么后来换成了 GAT。3.2 GAT 的注意力机制邻居不再一视同仁GAT 在聚合前引入注意力参数给不同的邻居边分配不同的权重。PyG 的GATConv封装了单头和多头注意力使用方式和GCNConv几乎一致你需要额外关注的只有两个参数heads和concat。from torch_geometric.nn import GATConv class GAT(nn.Module): def __init__(self, in_channels, hidden_channels, out_channels, heads8): super().__init__() # 第一层多头注意力输出通道数 hidden * heads self.conv1 GATConv(in_channels, hidden_channels, headsheads) # 第二层单头输出必须把维度降回 out_channels self.conv2 GATConv(heads * hidden_channels, out_channels, heads1) def forward(self, x, edge_index): x self.conv1(x, edge_index) x F.relu(x) x F.dropout(x, p0.6, trainingself.training) x self.conv2(x, edge_index) return x第一层如果设置heads8输出维度会变成hidden * 8第二层的输入通道就要相应调成heads * hidden_channels这是新手最常见的维度报错出处。另外 GAT 的原始论文在 Cora 上建议 dropout 为 0.6比 GCN 的 0.5 更高因为注意力机制更容易过拟合小而密的图。GAT 训练收敛通常比 GCN 慢在 Cora 上可能需要 300 到 400 epoch 才能稳定到 0.82 到 0.83 的准确率而 GCN 在 200 epoch 就完全收敛了。这不是代码问题而是注意力参数增加了模型容量。如果你的任务对推理延迟敏感先考虑 GCN 压缩版本如果更在意效果上限和邻居可解释性GAT 值得多等几分钟。3.3 GraphSAGE先采样后聚合大图唯一的可行路径GraphSAGE 解决的是另一个问题当图大到几亿条边、无法在显存里整图载入时GCN 和 GAT 都无能为力。GraphSAGE 的做法是每次迭代只采样固定数量的邻居一层采样 N 个两层就采样 N^2 个把计算限制在一个可控的局部子图上。PyG 里对应的类是SAGEConv它和 GCN 的最大区别是聚合后会用 拼接 的方式把中心节点特征和邻居特征合并而不是直接求和。from torch_geometric.nn import SAGEConv class GraphSAGE(nn.Module): def __init__(self, in_channels, hidden_channels, out_channels): super().__init__() self.conv1 SAGEConv(in_channels, hidden_channels) self.conv2 SAGEConv(hidden_channels, out_channels) def forward(self, x, edge_index): x self.conv1(x, edge_index) x F.relu(x) x F.dropout(x, p0.3, trainingself.training) x self.conv2(x, edge_index) return xSAGEConv本身不负责采样它只实现聚合逻辑。要真正跑 mini-batch 训练需要配合torch_geometric.loader.NeighborLoader按种子节点迭代采样邻域子图。常见做法是设置num_neighbors[25, 10]意思是一层采样 25 个邻居、二层采样 10 个采样参数直接决定内存占用和训练速度。3.4 一张表选型什么时候换模型参数从哪里调我不会让你把三种模型挨个试一遍再拍板那样代价太高。实际选型看三个信号图能不能整图放进显存、需不需要区分邻居重要程度、数据集的节点度分布是不是长尾。场景推荐模型关键参数说明小图整图训练、做基线实验GCNhidden 16lr 1e-2收敛最快代码最短图大放不进显存GraphSAGE NeighborLoadernum_neighbors[25, 10]用采样控制计算子图节点度差异大、边有噪声GATheads 8dropout 0.6注意力可以压制噪声边分类任务类别不平衡以上任意 macro F1不用默认 accuracy图数据简单准确率很容易虚高需要学多层结构GCN/GAT 控制在 2 到 3 层超过 3 层加残差连接防止过度平滑三个模型在 PyG 里共享完全相同的forward(x, edge_index)接口切换模型只需要替换model那一行。这种统一抽象是 PyG 最值得称赞的设计——模型选型阶段几乎不用改训练循环代码。实际项目中我常用的做法是先用 GCN 跑通全流程再在验证集上对比 GAT 的效果只有当 GAT 稳定高出 0.01 以上才值得换否则 GCN 的简单性本身就是收益。4. 训练一版完整的节点分类模型损失、mask 与超参调优4.1 训练集怎么切mask 是 GNN 半监督的底座图数据的训练集划分要谨慎。在普通机器学习里随机 7:1:2 划分是默认操作但好消息是 PyG 的内置数据集大多已经给你配好了官方划分。以 Cora 为例训练集 140 个节点、验证集 500、测试集 1000这不是随意切分的它是论文里约定的 benchmark。使用官方 mask 的好处是你可以读到公开论文里的准确率做横向对比——任何随机划分出来的结果都不具备可比较性。如果是自己的数据集必须遵循一个纪律划分的粒度是节点而不是子图。你要保证训练节点与测试节点虽然共享同一张图但训练节点不能拥有“看到测试节点邻居”的窗口。4.2 损失函数和优化器为什么是交叉熵和 AdamGNN 的传统做法是分类任务直接上交叉熵。原因有两个一是 PyG 的out输出的是每个节点在所有类别上的 logits交叉熵内部自带 softmax不会出现数值溢出问题二是 GNN 的训练集通常极小交叉熵在梯度传播上比 hinge loss 更容易把这个小数据集推到位。优化器上主流的两个选择是 Adam 和带动量的 SGD。GCN 论文里用的是 Adamlr0.01weight_decay5e-4。我实测过多次Adam 配上 1e-2 的学习率在 Cora 这类小图上表现极稳。def train(): model.train() optimizer.zero_grad() out model(data.x, data.edge_index) loss criterion(out[data.train_mask], data.y[data.train_mask]) loss.backward() optimizer.step() return float(loss)注意zero_grad()必须调用在loss.backward()之前否则梯度会累积到上一次迭代之上。图模型的参数量不大梯度累积问题在 CPU 上不容易察觉但如果后续换到 mini-batch 训练梯度累积造成的训练抖动会非常明显。4.3 五个必调参数hidden 维度、层数、lr、dropout、weight_decay先说 hidden 维度。Cora 这种 2708 节点的图hidden16 就够用了再往上加并不会带来明显的准确率提升反而让训练变慢。我实测 hidden 从 16 加到 64Cora 测试准确率浮动不超过 1 个点而训练时间翻了 4 倍。模型容量在图小的时候是过剩的——邻居聚合操作本身就是一种强正则化。层数上2 层是标准配置。上升到 3 层、4 层会有轻微提升但也引入过度平滑的风险。在商品推荐场景里3 层 GCN 能覆盖到“朋友的购买记录”这个远距离关联带来的关联能力提升比 hidden 翻倍更明显。层与层之间的 dropout 控制在 0.3 到 0.6规则很简单——模型泛化差就加大 dropout训练不收敛就减小。weight_decay 是 L2 正则GCN 原始参数 5e-4 可以用但如果你用的是大图建议调小到 1e-5 到 1e-4 区间。4.4 训练监控损失下降但验证不动问题在 mask 不在模型我建议每个实验都同时打印训练集、验证集两条损失曲线。GNN 里最常见的训练陷阱是训练损失一路下降、验证损失纹丝不动——这通常不是过拟合而是你的验证集里混入了训练节点的邻居信息验证节点被“剧透”了。解决方式是回到数据切分阶段检查验证集节点是否与训练集节点直接相连或者只是在model.eval()时忘记关闭 dropout。另外需要在训练循环里保存一份验证集最优的模型权重否则最后拿到的是最后一个 epoch 的权重而这个 epoch 可能已经在验证集上退化了。best_val_acc 0 best_state None for epoch in range(200): loss train() val_acc evaluate(val_mask) if val_acc best_val_acc: best_val_acc val_acc best_state {k: v.clone() for k, v in model.state_dict().items()} if epoch % 20 0: print(fEpoch {epoch}, Loss {loss:.4f}, Val {val_acc:.4f}) model.load_state_dict(best_state) test_acc evaluate(test_mask)这招可以说是 GNN 训练的后悔药。最终测试结果以best_state为准而不是最后一轮的权重。图模型收敛曲线比普通 DNN 更震荡中等规模的图上 10 个 epoch 的波动就能吃掉一个点准确率。这套代码逻辑上等效于早停但又不用定义窗口大小比 early stopping 更适合小数据集。5. GNN 代码复现的排障与避坑记录5.1 现象import torch_geometric 报错numpy 版本冲突原因PyG 对底层的 torch、numpy 版本非常敏感是 GNN 环境里最常见的坑。初次安装时用 pip 直接安装依赖解析器经常选中一个过新的 numpy 版本和 torch 自带的版本对不上import 时直接报错或者报AttributeError: module numpy has no attribute int。解决严格按官网给的安装指令来并且安装完在 Python 里 import 一下验证。如果出现 numpy 冲突通常是 numpy 版本高于 torch 预期版本把 numpy 降级到 torch 官方要求的区间内。建议创建一个专门的 conda 环境不要用系统级 Python 混合安装不同项目共用环境是 GNN 开发最大的时间黑洞。5.2 现象测试准确率只有 0.55和论文 0.81 差距巨大原因八成是你用了random_split或者自己做了一层train_test_split破坏了图数据的半监督结构。你随机抽出来的训练集节点和测试集节点直接相邻训练损失反传时网络“见过”测试节点附近的信息但 mask 的规模又太小导致模型学不全。另外一个低概率原因是数据集下载成了不完整的子集某些版本缓存了损坏的.pt文件。解决先确认data.train_mask.sum()是否等于 140 而不是别的数。然后改用官方 maskdata.train_mask/data.val_mask/data.test_mask不要自己造。如果必须用自定义划分请先跑一个节点不重叠的分层子图拆分脚本或者干脆用标签传播的方式先看一下原始图的同质性。最稳妥的做法是在 PyG 内置数据集上先复现论文指标再迁移到你自己的数据上。5.3 现象边数异常模型震荡不收敛原因edge_index里存在大量重复边或者数据集的边没有转成无向。自己从 CSV 构造数据时很多人在去重阶段把(u, v)和(v, u)当成重复对处理结果整个消息传播方向被切断。还有一类原因是边列表里包含了自环同时模型内部又加了add_self_loops导致自环被重复计数两次。解决统一的做法是用 PyG 的工具函数清洗边在进入模型前跑一遍to_undirected和coalesce这样能同时处理方向、去重和排序。from torch_geometric.utils import to_undirected, coalesce edge_index coalesce(edge_index) edge_index to_undirected(edge_index)处理完打印edge_index[0].unique().shape确认所有节点都在边列表里出现过。孤立的节点没有邻居消息传递对它不更新它的预测就只能靠自身特征很容易把整体准确率拉成一个奇怪的形状。5.4 现象显存溢出但图明明不算大原因你实际上做的是“全批度图传播”而不是“批量采样”。GCN 在forward时会把整张图的邻接矩阵转成密集形式参与运算如果邻接矩阵是 10 万乘 10 万光是矩阵本身就需要不可忽略的显存。PyG 内部通常会自动把稀疏邻接矩阵放到稀疏张量里但如果你手动在 CPU 上执行了torch.tensor转换就会导致显存瞬间翻倍。解决让edge_index保持 long 型稀疏格式不要手动调用.to_dense()。如果图实在太大用NeighborLoader做邻居采样把计算限制在子图上同时把批量数据移动到 GPU而不是提前把整图搬进去。5.5 现象训练时准确率极高评估时直接掉到随机水平原因训练和评估模式没有正确切换。model.eval()只切换了 dropout 和 BN 的行为但如果你没把 mask 从数据里取出来评估时会把所有节点的 loss 都算上包括大量无标签节点。无标签节点有随机的 label自然会把评价指标拉低。另一个隐藏原因是数据里存在缺失的 target在数据加载时用fillna或者直接滤掉未过滤的节点带上了默认标签被算进了评估集合。解决训练用out[data.train_mask]评估用out[data.test_mask]两处都单独写不要图省事用同一个 out 变量。一个更稳妥的做法是封装一个evaluate(mask, model, data)函数把 mask 作为参数传递避免评估时手抖用了 train_mask。6. 从节点分类走到链路预测邻居聚合的进阶方向与工程验证节点分类跑通只是第一步。在真实系统里比如推荐系统的“猜你喜欢”你真正需要的往往是说清“两个物品有没有可能发生交互”这是链路预测任务。实现链路预测的基础正是你已经写好的 GNN encoder——用 GCN 或 GAT 学习出每个节点的向量表示然后把两两向量的相似度映射成一个连接概率。PyG 里做这件事有标准框架把图的边拆成训练边和验证边训练时用train_test_split_edges()把边分成 positive 和 negative 样本再用torch_geometric.nn.GAE包装你的 GNN encoder。值得注意的技巧是链路预测的负样本不存在的边必须在训练时采样而不是拿原图说了算——你和邻居之间的负样本不够模型会把所有节点都预测成“有边”。工程上的验证方法也不一样。链路预测的主要评估指标是 AUC 和 HitsKAUC 直接反映排序质量Hits100 在推荐场景里更有实际意义。我第一次做链路预测时只看准确率结果模型学出来对稀疏边和稠密边的偏好完全不一样后来改成 AUC Hits 双指标才把问题暴露出来。顺便提一个不容易注意的地方图神经网络在新节点增量上线时不重新训练也能直接推理新节点——前提是它的邻居特征已经存在。这做线上服务的人应该特别留意因为这意味着你可以用训练好的 GAT 去推断冷启动用户只要他至少有一条关系边接口。我自己现在的 GNN 工程管线也慢慢形成了一套定式所有模型的edge_index都预处理成无向去重格式数据集划分统一用节点 mask 而不是随机切训练过程只保留验证集最优的权重。遵照这三条习惯之后踩过的坑明显少了很多。希望这些代码和参数能帮你在相似的方向少走一点弯路。本文还有配套的精品资源点击获取