简介:这份资源面向希望入门或进阶图神经网络预测任务的开发者与算法学习者,提供一套可直接运行的Python完整源码数据包,帮助解决从模型搭建到结果复现的实践问题。压缩包共32个文件,约8.34MB,以19个py脚本和4个ipynb交互式笔记为核心,辅以4个npz数据文件、requirements.txt依赖清单、setup.py安装脚本及README说明,覆盖数据预处理、模型定义与训练流程等模块。内容围绕PPNP等图神经网络方法展开,包含PyTorch与TensorFlow两套实现示例及结果复现笔记,便于读者对照不同框架理解消息传递与预测机制。目前已有2252人学习下载,适合作为课程设计、科研实验或工程原型的参考起点,读者可借此快速掌握图神经网络预测的代码组织方式与实验复现路径。
1. 从「跑通一个 GNN 预测脚本」说起:这套 Python 源码包到底解决什么问题
很多人第一次接触图神经网络,卡住的地方不是公式,而是数据怎么变成图、模型怎么跑起来、预测结果怎么落回业务。标题里说的「基于 GNN 图神经网络预测(Python 完整源码数据包)」,本质是一套能直接跑的端到端方案:把结构化或关系型数据构造成图,用 GNN 做节点、边或整图级别的预测,最后输出可用的预测值。它适合两类人——一类是想快速验证 GNN 在自己数据上有没有效果的算法工程师,另一类是有 Python 基础、想从 sklearn 那套表格模型跨到图模型的数据分析从业者。我见过太多人把 PyG、DGL 装完就停在 import 报错,或者数据一换就维度对不上,所以这篇不聊虚的,直接按「数据构图 → 模型定义 → 训练 → 预测 → 排错」把可复现的路径写清楚,参数怎么调、坑在哪,都落到代码上。
2. GNN 预测任务拆解:先分清你是节点预测、边预测还是图预测
2.1 三类预测任务对应的数据形态和输出
GNN 不是一种模型,而是一类在图上做消息传递的框架。落到预测任务上,常见分三种:节点级预测(给每个节点打标签,比如用户是否流失)、边级预测(判断两个节点之间有没有关系,比如推荐里的点击概率)、图级预测(整张图对应一个标签,比如分子属性预测)。这三种任务的数据组织方式完全不同,选错任务类型,后面代码全白写。
节点预测的输入是一张图加每个节点的特征矩阵,输出是节点数 × 类别数。边预测要把边拆成正负样本,输出是边上的二分类或多分类。图级预测则需要一批独立的图,每张图一个标签。我一般建议新手从节点预测入手,因为数据构造最简单,调试链路最短。
| 任务类型 | 输入形态 | 输出维度 | 典型场景 |
|---|---|---|---|
| 节点预测 | 单图 + 节点特征 | N × C | 用户分类、欺诈检测 |
| 边预测 | 单图 + 边样本 | E × C | 推荐、知识图谱补全 |
| 图预测 | 图列表 + 图标签 | B × C | 分子性质、社区分类 |
2.2 用 PyTorch Geometric 构造第一张图
环境上,Python 3.9 到 3.11 都能跑,PyG 对 torch 版本有绑定关系,装之前先确认 torch 版本。下面这段是最小可跑的构图代码,把一份边表和节点特征转成 PyG 的 Data 对象。
import torch from torch_geometric.data import Data # edge_index 必须是 [2, E] 的 long 张量,第一行是源节点,第二行是目标节点 edge_index = torch.tensor([ [0, 1, 1, 2, 3, 4], [1, 0, 2, 1, 4, 3] ], dtype=torch.long) # 节点特征矩阵,N 个节点,每个节点 F 维特征 x = torch.tensor([ [1.0, 0.2], [0.5, 0.8], [0.3, 0.1], [0.9, 0.6], [0.2, 0.4] ], dtype=torch.float) # 节点标签,这里假设 3 分类 y = torch.tensor([0, 1, 2, 1, 0], dtype=torch.long) data = Data(x=x, edge_index=edge_index, y=y) print(data)这段代码里,edge_index的方向性很关键。如果你的图是无向图,每条边要写两遍,正反各一次,否则消息传递只能单向流动,模型效果会明显偏弱。x的每一行对应一个节点,行号就是节点 ID,必须和edge_index里的编号严格对应。y的长度必须等于节点数,少一个都会在算 loss 时报维度错误。
参数上,dtype一定要显式指定,torch.long用于索引和标签,torch.float用于特征。很多人从 pandas 转过来忘了转 dtype,结果 PyG 内部报expected Long but got Float,这种错排查起来很费时间。
2.3 划分训练、验证、测试集的正确姿势
图数据不能像表格数据那样随机切分,因为节点之间有关系,随机切会导致信息泄漏。常见做法是用掩码,把节点分成三部分,训练时只回传训练节点的梯度。
from torch_geometric.transforms import RandomNodeSplit # 按 60/20/20 划分,num_classes 用于分层采样 split = RandomNodeSplit(num_val=0.2, num_test=0.2) data = split(data) print(data.train_mask.sum(), data.val_mask.sum(), data.test_mask.sum())RandomNodeSplit会生成三个布尔掩码,长度等于节点数。训练时用train_mask选节点算 loss,验证和测试同理。如果你的图是多个独立小图,就不能用节点掩码,要按图切分,这时候用DataLoader配合图列表。这里有个容易翻车的点:掩码划分后如果某类节点在训练集里一个都没有,模型直接学不会这一类,所以分类任务最好用分层划分,或者手动检查每类在训练集中的占比。
3. 模型定义与训练:从两层 GCN 到一个能收敛的预测器
3.1 为什么先上 GCN,什么时候换 GAT 或 GraphSAGE
GCN 是最基础的图卷积,本质是对邻居特征做加权平均再线性变换。它的假设是邻居重要性相同,实现简单,收敛稳定,适合作为第一个 baseline。GAT 引入注意力,让模型自己学邻居权重,在异质图或邻居质量差异大的场景下更强,但参数量上去后小数据集容易过拟合。GraphSAGE 适合大图,因为它对邻居采样,不用一次性加载全图。
我一般会先用两层 GCN 跑通,看验证集指标。如果 GCN 明显欠拟合,再换 GAT;如果图特别大、显存吃紧,换 GraphSAGE 加邻居采样。选型不是越新越好,而是看你的图规模和任务难度。
3.2 两层 GCN 的完整定义与维度对齐
import torch.nn.functional as F from torch_geometric.nn import GCNConv class GCNNet(torch.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): x = self.conv1(x, edge_index) x = F.relu(x) x = F.dropout(x, p=0.5, training=self.training) x = self.conv2(x, edge_index) return xin_channels必须等于节点特征的列数,out_channels等于类别数。隐藏层维度一般取 16 到 128,图越大可以适当加宽。dropout放在两层之间,训练时生效,推理时关闭,这是防止过拟合最直接的手段。如果验证集 loss 一直震荡,先把 dropout 调到 0.3 试试。
3.3 训练循环与关键超参数
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = GCNNet(in_channels=data.num_features, hidden_channels=32, out_channels=3).to(device) data = data.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4) for epoch in range(200): model.train() optimizer.zero_grad() out = model(data.x, data.edge_index) loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask]) loss.backward() optimizer.step() if epoch % 20 == 0: model.eval() with torch.no_grad(): pred = model(data.x, data.edge_index).argmax(dim=1) val_acc = (pred[data.val_mask] == data.y[data.val_mask]).float().mean() print(f'Epoch {epoch}, Loss {loss:.4f}, Val Acc {val_acc:.4f}')学习率 0.01 是 GCN 的常用起点,weight_decay用 5e-4 做 L2 正则。训练轮数看验证集指标,一般 100 到 500 轮之间。如果 loss 下降但验证准确率不涨,说明过拟合,加 dropout 或减隐藏层维度。如果 loss 不降,先检查edge_index有没有写反、标签有没有对齐。
4. 预测输出与结果落盘:把模型输出变成能用的数据
4.1 推理阶段关闭梯度并导出预测
训练完不能直接拿model()的结果当最终输出,要切到 eval 模式并关梯度,否则 dropout 和 batchnorm 行为不一致,预测结果会抖。
model.eval() with torch.no_grad(): logits = model(data.x, data.edge_index) probs = F.softmax(logits, dim=1) preds = probs.argmax(dim=1) # 导出为 numpy 再写文件 import numpy as np result = np.stack([preds.cpu().numpy(), probs.max(dim=1).values.cpu().numpy()], axis=1) np.savetxt('gnn_predictions.csv', result, delimiter=',', header='label,confidence', comments='')probs.max(dim=1).values是每个节点预测的置信度,落盘时一起带上,方便后续按阈值过滤。np.savetxt的comments=''是为了不让 header 前面多一个#,方便直接给下游读。
4.2 用 pandas 做结果后处理和写入 Excel
很多业务方要的是 Excel,不是 csv。用 pandas 转一手最省事。
import pandas as pd df = pd.DataFrame(result, columns=['label', 'confidence']) df['node_id'] = df.index # 置信度低于 0.6 的标记为待人工复核 df['need_review'] = df['confidence'] < 0.6 df.to_excel('gnn_predictions.xlsx', index=False)need_review这一列在实际业务里很有用,模型不是万能的,低置信度的样本交给人工,比硬信模型输出靠谱。写 Excel 需要openpyxl,没装的话pip install openpyxl就行。
5. 避坑与排查:GNN 预测里最容易翻车的 5 个地方
5.1 现象:loss 变成 nan,训练几轮就崩
原因通常是学习率太大或者特征没归一化。GCN 对特征尺度敏感,如果某一维特征数值到几千,梯度直接爆炸。解决方法是先对x做标准化,x = (x - x.mean(0)) / (x.std(0) + 1e-6),再把学习率降到 0.001 试。
5.2 现象:验证集准确率一直不动,和随机猜差不多
先查edge_index方向。如果无向图只写了一遍边,消息传不回去,节点只能看到一部分邻居。再查标签和节点 ID 是否错位,pandas 的 index 和 PyG 的节点编号经常对不上,建议构图前先reset_index(drop=True)。
5.3 现象:换一份数据就报维度不匹配
in_channels写死了,新数据特征列数变了。正确做法是从data.num_features动态取,别硬编码。同理out_channels从data.y.max() + 1或len(classes)取。
5.4 现象:GPU 显存不够,大图跑不起来
全图训练对显存要求高。常见做法是转成邻居采样训练,用NeighborLoader,每次只加载一批节点的子图。batch_size 从 512 开始调,显存还是不够就降到 128。
5.5 现象:预测结果全是同一类
类别极度不平衡时,模型会偷懒全预测多数类。解决是在 loss 里加weight参数,按类别频率的倒数给权重,或者对少数类做上采样。先打印训练集每类的数量,心里有数再调。
6. 进阶技巧:让 GNN 预测从「能跑」到「敢用」
模型能跑通只是第一步,真正上线前我会做两件事:一是用早停(early stopping)防止过拟合,二是做一次特征重要性或注意力权重的可视化,确认模型没学到噪声。
早停的实现很简单,在验证集上监控指标,连续 N 轮不提升就停。
best_val = 0 patience = 20 counter = 0 for epoch in range(500): # ... 训练代码 ... if val_acc > best_val: best_val = val_acc counter = 0 torch.save(model.state_dict(), 'best_model.pt') else: counter += 1 if counter >= patience: print(f'Early stop at epoch {epoch}') breakpatience一般设 10 到 30,太小容易停早,太大浪费算力。保存最优权重比保存最后一轮权重更稳,因为最后一轮可能已经过拟合了。
另一个技巧是给边加权重。如果你的图里不同边的重要性不同,可以在GCNConv里传edge_weight,让消息传递时按权重聚合。这在交易网络、社交网络里很常见,权重可以用交易金额、互动频率来定。
最后说个我自己的习惯:每次换数据集,先画一下度分布。如果大部分节点度是 1 或 2,GNN 的优势根本发挥不出来,这时候用 MLP 加节点特征可能效果差不多,别硬上图模型。GNN 的价值在于邻居信息,图太稀疏就失去意义了。这套源码包的价值不在于代码多复杂,而在于把构图、训练、预测、落盘这条链路串通了,你换自己的数据时只需要改数据加载那一段。希望帮到你。
本文还有配套的精品资源,点击获取