资讯动态

PyTorch Geometric节点分类实战:从GCN到GAT与GraphSAGE

发布时间:2026/10/17 16:20:31 来源:尧图企业网站定制
1. 从零构建节点分类任务为什么选PyTorch Geometric搞图神经网络的人绕不开节点分类这个任务。它算是GNN领域的“Hello World”但真正动手做的时候你会发现从数据加载到模型训练中间有一堆琐碎的工程细节。我最初用原生PyTorch手撸GCN光是稀疏邻接矩阵的构建和归一化就写了快两百行还容易出错。后来换成PyTorch Geometric简称PyG同样的工作压缩到十几行而且内置了大量图卷积层和数据集接口省下来的时间可以专注在模型设计和调参上。这篇文章面向的是已经了解PyTorch基础、想快速上手图神经网络节点分类的开发者。不管你是做社交网络分析、引用网络分类还是分子属性预测节点分类的流程都是相通的。我会从数据准备、模型搭建、训练循环、结果评估到常见坑点完整走一遍。代码可以直接复现参数选择我也会解释清楚背后的逻辑。PyG的核心优势在于它把图数据抽象成了Data对象节点特征、边索引、标签都封装在一起配合DataLoader可以方便地做mini-batch训练。另外它提供了GCNConv、GATConv、SAGEConv等常用卷积层直接调用即可。对于节点分类最经典的三个数据集是Cora、Citeseer和Pubmed都是引用网络节点是论文边是引用关系特征是词袋向量标签是论文类别。下面我就以Cora为例展开。2. 环境搭建与PyG安装避坑指南2.1 版本匹配是最大的坑PyG的安装不像普通pip包那么直接因为它依赖torch-scatter、torch-sparse等扩展库这些库需要和PyTorch版本、CUDA版本严格对应。我踩过最惨的一次是PyTorch 1.12配了torch-scatter 2.1直接报符号未定义错误。所以第一步永远是确认版本。先查你的PyTorch版本和CUDA版本import torch print(torch.__version__) print(torch.version.cuda)假设输出是2.0.1和11.8那么安装命令应该这样写pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-2.0.1cu118.html pip install torch-geometric如果你没有GPU或者不想折腾CUDA直接用CPU版本也行把URL里的cu118换成cpu即可。注意PyG的轮子文件托管在data.pyg.org国内下载可能慢可以加个镜像或者手动下载whl文件再pip install。提示不要用conda install pytorch-geometricconda源里的版本往往滞后而且依赖解析容易出问题。pip方式最稳。2.2 验证安装是否成功装完之后跑一段小代码验证import torch from torch_geometric.data import Data edge_index torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtypetorch.long) x torch.tensor([[-1], [0], [1]], dtypetorch.float) data Data(xx, edge_indexedge_index) print(data)如果输出类似Data(edge_index[2, 4], x[3, 1])说明环境没问题。如果报ImportError大概率是torch-scatter没装好回头检查版本对应关系。2.3 数据集下载与缓存PyG内置了Planetoid数据集接口Cora、Citeseer、Pubmed都能一键下载from torch_geometric.datasets import Planetoid dataset Planetoid(root/tmp/Cora, nameCora)第一次运行会自动下载到/tmp/Cora目录。如果网络不通可以手动下载cora.zip放到对应目录下。数据集包含140个训练节点、500个验证节点、1000个测试节点特征维度1433类别数7。这个划分是固定的方便大家公平对比。3. 数据对象拆解与核心属性解析3.1 Data对象里到底有什么加载完数据集后dataset[0]就是一个Data对象。我们逐项看它的属性data dataset[0] print(data) print(data.x.shape) # [2708, 1433] print(data.edge_index.shape) # [2, 10556] print(data.y.shape) # [2708] print(data.train_mask.sum()) # 140x节点特征矩阵2708个节点每个节点1433维词袋特征。edge_index边索引形状[2, E]第一行是源节点第二行是目标节点。注意这是有向图表示但Cora通常按无向图处理所以每条边会出现两次正反各一次。y节点标签取值0到6。train_mask、val_mask、test_mask布尔掩码标记哪些节点用于训练、验证、测试。3.2 特征归一化要不要做Cora的特征已经是词袋向量数值范围0到1左右通常不需要额外归一化。但如果你换成其他数据集比如分子图特征可能跨度很大这时候做行归一化L1或L2会有帮助。我试过在Cora上做L2归一化准确率变化不到0.5%所以默认不做也行。3.3 边索引的对称化处理虽然Cora的edge_index已经包含了双向边但如果你自己构建图记得手动对称化from torch_geometric.utils import to_undirected edge_index to_undirected(edge_index)否则GCN在聚合邻居信息时只能看到单向邻居效果会打折扣。这个细节很多教程没提但实际影响不小。4. GCN模型搭建从公式到代码4.1 GCN的核心公式回顾GCN的层间传播公式是[ H^{(l1)} \sigma(\tilde{D}^{-1/2} \tilde{A} \tilde{D}^{-1/2} H^{(l)} W^{(l)}) ]其中(\tilde{A} A I)是加了自环的邻接矩阵(\tilde{D})是(\tilde{A})的度矩阵(W)是可学习权重(\sigma)是激活函数。PyG的GCNConv已经内置了归一化逻辑我们只需要定义权重和激活。4.2 两层GCN的代码实现import torch import torch.nn.functional as F from torch_geometric.nn import GCNConv class GCN(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, p0.5, trainingself.training) x self.conv2(x, edge_index) return F.log_softmax(x, dim1)这里隐藏层维度设为16输出维度为7类别数。Dropout率0.5是论文里的默认值实测下来0.3到0.6之间差别不大但太低容易过拟合太高欠拟合。4.3 为什么用log_softmax而不是softmax因为训练时用F.nll_loss负对数似然损失它要求输入是对数概率。如果你用CrossEntropyLoss那模型输出原始logits即可不需要log_softmax。两者等价看个人习惯。我一般用nll_loss配log_softmax因为PyG官方示例就是这么写的。4.4 权重初始化的小技巧PyG的卷积层默认用Xavier初始化但你可以手动改def init_weights(m): if isinstance(m, GCNConv): torch.nn.init.xavier_uniform_(m.lin.weight) if m.lin.bias is not None: torch.nn.init.zeros_(m.lin.bias)我试过用Kaiming初始化在Cora上准确率差不多但在更深的网络上Xavier更稳。所以默认就好不用折腾。5. 训练循环与评估指标实战5.1 标准训练流程device torch.device(cuda if torch.cuda.is_available() else cpu) model GCN(dataset.num_features, 16, dataset.num_classes).to(device) data dataset[0].to(device) optimizer torch.optim.Adam(model.parameters(), lr0.01, weight_decay5e-4) def train(): model.train() optimizer.zero_grad() out model(data.x, data.edge_index) loss F.nll_loss(out[data.train_mask], data.y[data.train_mask]) loss.backward() optimizer.step() return loss.item() def test(): model.eval() out model(data.x, data.edge_index) pred out.argmax(dim1) accs [] for mask in [data.train_mask, data.val_mask, data.test_mask]: correct (pred[mask] data.y[mask]).sum() accs.append(int(correct) / int(mask.sum())) return accs训练200个epoch每10个epoch打印一次损失和准确率。Cora上通常能达到81%到83%的测试准确率。5.2 学习率与权重衰减的选择学习率0.01、权重衰减5e-4是GCN论文里的配置。我做过一组对比实验学习率权重衰减测试准确率0.015e-481.5%0.011e-380.8%0.0055e-481.2%0.025e-480.1%可以看到默认配置已经接近最优。权重衰减再大容易欠拟合学习率再大训练不稳定。5.3 早停策略的实现验证集准确率连续50个epoch不提升就停止训练best_val_acc 0 patience 50 counter 0 for epoch in range(1, 500): loss train() train_acc, val_acc, test_acc test() if val_acc best_val_acc: best_val_acc val_acc counter 0 torch.save(model.state_dict(), best_model.pt) else: counter 1 if counter patience: print(fEarly stopping at epoch {epoch}) break早停能防止过拟合尤其在数据量小的时候效果明显。Cora只有140个训练节点不加早停的话测试准确率可能掉1到2个百分点。6. 从GCN扩展到GAT与GraphSAGE6.1 GAT的注意力机制GAT用注意力权重代替固定的归一化系数公式里多了个注意力系数(\alpha_{ij})。PyG的GATConv可以直接替换GCNConvfrom torch_geometric.nn import GATConv class GAT(torch.nn.Module): def __init__(self, in_channels, hidden_channels, out_channels, heads8): super().__init__() self.conv1 GATConv(in_channels, hidden_channels, headsheads) self.conv2 GATConv(hidden_channels * heads, out_channels, heads1) def forward(self, x, edge_index): x F.dropout(x, p0.6, trainingself.training) x F.elu(self.conv1(x, edge_index)) x F.dropout(x, p0.6, trainingself.training) x self.conv2(x, edge_index) return F.log_softmax(x, dim1)注意第一层用了8个注意力头输出维度是hidden_channels * heads第二层用1个头输出类别数。GAT在Cora上能到83%左右比GCN略高但训练更慢显存占用更大。6.2 GraphSAGE的邻居采样GraphSAGE适合大图因为它对每个节点采样固定数量的邻居避免全图聚合。PyG的SAGEConv支持这种模式from torch_geometric.nn import SAGEConv class GraphSAGE(torch.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 F.relu(self.conv1(x, edge_index)) x F.dropout(x, p0.5, trainingself.training) x self.conv2(x, edge_index) return F.log_softmax(x, dim1)在Cora这种小图上GraphSAGE和GCN表现接近但在大规模图比如Reddit上GraphSAGE的采样策略能显著降低显存和计算量。6.3 三种模型的对比模型测试准确率训练时间/epoch显存占用GCN81.5%0.8s低GAT83.2%2.1s中GraphSAGE81.8%1.0s低选择哪个取决于你的需求追求精度选GAT追求速度选GCN大图选GraphSAGE。7. 常见报错与排查技巧实录7.1 CUDA out of memoryCora很小一般不会OOM。但如果你换成Pubmed或自己构建的大图显存可能不够。解决方法减小隐藏层维度比如从64降到16。用torch.cuda.empty_cache()清理缓存。改用GraphSAGE的邻居采样避免全图聚合。7.2 准确率不升反降训练过程中验证集准确率波动很大甚至下降常见原因学习率太大尝试降到0.001。Dropout率太低过拟合了调到0.5以上。没有做早停模型在训练集上过拟合。我遇到过一次准确率卡在70%上不去后来发现是edge_index没有对称化补上to_undirected之后直接跳到81%。7.3 损失为NaN通常是学习率过大或者特征里有异常值。检查data.x是否有inf或nanprint(torch.isnan(data.x).sum()) print(torch.isinf(data.x).sum())如果有用torch.nan_to_num处理。另外梯度裁剪也能缓解torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)7.4 验证集准确率高于测试集这很正常因为验证集参与了早停选择相当于间接调参。如果差距超过5个百分点说明验证集太小或者分布不一致。Cora的验证集有500个节点差距通常在1到2个百分点。8. 调参经验与性能优化心得8.1 隐藏层维度不是越大越好我试过16、32、64、128四种维度隐藏层维度测试准确率参数量1681.5%23K3281.7%46K6481.3%92K12880.9%184K16到32差别不大再大反而过拟合。Cora只有140个训练节点模型容量太大容易记住训练集。8.2 层数不宜过深两层GCN效果最好三层以上会出现过平滑over-smoothing所有节点特征趋同准确率下降。我试过4层GCN测试准确率掉到75%左右。如果非要加深可以加残差连接x x self.conv1(x, edge_index)8.3 特征归一化的影响前面提过Cora不需要归一化但如果你用其他数据集比如蛋白质图特征跨度大做标准化减均值除标准差能提升5%以上。这个要根据数据分布决定。8.4 随机种子的影响GNN对随机初始化敏感不同种子下准确率可能差1到2个百分点。建议跑5次取平均import numpy as np accs [] for seed in range(5): torch.manual_seed(seed) np.random.seed(seed) # 训练和测试 accs.append(test_acc) print(fAverage accuracy: {np.mean(accs):.4f} ± {np.std(accs):.4f})这样报告的结果更可靠。9. 完整代码整合与运行说明把前面的片段拼起来形成一个可运行的脚本import torch import torch.nn.functional as F from torch_geometric.datasets import Planetoid from torch_geometric.nn import GCNConv # 加载数据 dataset Planetoid(root/tmp/Cora, nameCora) data dataset[0] # 定义模型 class GCN(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 F.relu(self.conv1(x, edge_index)) x F.dropout(x, p0.5, trainingself.training) x self.conv2(x, edge_index) return F.log_softmax(x, dim1) # 训练配置 device torch.device(cuda if torch.cuda.is_available() else cpu) model GCN(dataset.num_features, 16, dataset.num_classes).to(device) data data.to(device) optimizer torch.optim.Adam(model.parameters(), lr0.01, weight_decay5e-4) # 训练循环 best_val_acc 0 for epoch in range(1, 201): model.train() optimizer.zero_grad() out model(data.x, data.edge_index) loss F.nll_loss(out[data.train_mask], data.y[data.train_mask]) loss.backward() optimizer.step() model.eval() pred model(data.x, data.edge_index).argmax(dim1) val_acc (pred[data.val_mask] data.y[data.val_mask]).float().mean() if val_acc best_val_acc: best_val_acc val_acc torch.save(model.state_dict(), best_model.pt) if epoch % 20 0: test_acc (pred[data.test_mask] data.y[data.test_mask]).float().mean() print(fEpoch {epoch:03d}, Loss: {loss:.4f}, Val: {val_acc:.4f}, Test: {test_acc:.4f}) # 加载最佳模型测试 model.load_state_dict(torch.load(best_model.pt)) model.eval() pred model(data.x, data.edge_index).argmax(dim1) test_acc (pred[data.test_mask] data.y[data.test_mask]).float().mean() print(fFinal Test Accuracy: {test_acc:.4f})运行这个脚本你应该能看到测试准确率在0.81到0.83之间。如果低于0.80检查一下PyG版本和随机种子。10. 后续扩展方向与个人建议节点分类只是图神经网络的入门任务掌握之后可以往几个方向深入。一是异构图比如学术网络里有作者、论文、会议三种节点用HeteroConv处理。二是动态图边和节点随时间变化需要时序图神经网络。三是大规模图用NeighborLoader做mini-batch训练避免全图加载。我个人在实际项目里用得最多的是GraphSAGE加邻居采样因为工业级图往往有百万甚至千万节点全图训练不现实。另外特征工程也很关键原始特征不好的话再复杂的模型也救不回来。我一般会先跑一个逻辑回归基线如果GNN比基线高不到3个百分点说明图结构带来的增益有限得回头检查图构建是否合理。最后分享一个小技巧调试GNN时先把模型调到极简比如一层GCN、隐藏层8维、不加Dropout看能否在训练集上过拟合。如果连训练集都拟合不了说明代码有bug而不是调参问题。这个思路帮我省了很多时间。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑