PyTorch Geometric 实战教程：10 行代码训练你的第一个图神经网络

PyTorch Geometric 实战教程10 行代码训练你的第一个图神经网络【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometricPyTorch GeometricPyG是 PyTorch 生态中事实标准的关系型深度学习库它把图、点云、分子等非规整结构数据统一成张量再配上消息传递框架与上百个预实现 GNN 层让你用写普通 PyTorch 模型的方式训练图神经网络。读完本文你可以三步完成安装与最小可运行的 GCN 训练首次运行即可看到 Cora 数据集上的准确率理解 PyG 的三大核心机制MessagePassing消息传递、EdgeConv动态建图、PointTransformerConv点云注意力并能定位到仓库源码跑通一个端到端的 3D 点云分类案例DGCNN on ModelNet10知道如何验证结果掌握邻居采样小批量训练与多 GPU 分布式训练两种扩展手段图数据为什么难训练PyG 如何解决用传统卷积神经网络训练图像时输入永远是规整的二维网格。但真实世界的很多数据并不满足这一点常见的痛点有数据结构不规则无法直接喂给 CNN引用网络、社交网络、分子、3D 点云的邻域关系各不相同节点数量、边数量每批数据都不一样。PyG 用edge_index形状[2, num_edges]的边索引张量描述任意稀疏连接关系卷积操作退化为一次按边聚合torch_geometric/data/中的Data对象统一管理节点特征、边、标签与批划分。从零实现 GNN 层成本很高手写一次消息传递需要处理边的索引、邻居聚合、批次边界等细节。PyG 在torch_geometric/nn/conv/中实现了 60 多种标准层GCNConv、GATConv、SAGEConv、PointNetConv……全部继承自统一的MessagePassing接口且支持torch.compile与 TorchScript。单机放不下大图亿级节点的图无法整图放进 GPU。PyG 提供NeighborLoader、ClusterLoader等采样式加载器见torch_geometric/loader/以及基于torch.distributed的多 GPU 训练路径见 examples/multi_gpu/。三步跑通开发环境第一步安装。从 PyG 2.3 起核心功能只依赖 PyTorch 本身pip install torch_geometric # 可选安装加速扩展稀疏算子加速按需选择 CUDA 版本 pip install pyg_lib torch_scatter torch_sparse \ -f https://data.pyg.org/whl/torch-2.10.0cu126.html官方建议先做最小安装需要SplineConv、异构图算子等特性时再补装扩展。第二步写最小模型。下面是 PyG 官方 README 中的快速示例加载 Cora 引用网络堆叠两层GCNConvfrom torch_geometric.nn import GCNConv from torch_geometric.datasets import Planetoid dataset Planetoid(root., nameCora) # 自动下载并解析引用图 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).relu() # 一次消息传递按边聚合邻居 return self.conv2(x, edge_index) model GCN(dataset.num_features, 16, dataset.num_classes)第三步训练循环。与标准 PyTorch 完全一致唯一的区别是损失只在有标签的节点上计算data dataset[0] # Data 对象x 为节点特征edge_index 为边索引 optimizer torch.optim.Adam(model.parameters(), lr0.01) for epoch in range(200): pred model(data.x, data.edge_index) # 半监督设定只在训练掩码节点上计算交叉熵 loss torch.nn.functional.cross_entropy(pred[data.train_mask], data.y[data.train_mask]) optimizer.zero_grad() loss.backward() optimizer.step()运行后你会看到 Cora 上约 81% 的测试准确率完整评估代码见 examples/gcn.py。数据集加载、T.NormalizeScale()、T.SamplePoints(1024)等变换的使用方式可参考 examples/dgcnn_classification.py 的 42–86 行。核心机制剖析消息传递、动态建图与点云注意力一切卷积层都长在 MessagePassing 上GCNConv的forward只做了三件事源码见 torch_geometric/nn/conv/gcn_conv.pydef forward(self, x, edge_index, edge_weightNone): if self.normalize: # gcn_norm 补自环并计算对称归一化权重 D^{-1/2} A D^{-1/2} edge_index, edge_weight gcn_norm(edge_index, edge_weight, x.size(self.node_dim), self.improved, self.add_self_loops, self.flow, x.dtype) x self.lin(x) # 特征变换 X Theta # propagate 内部调用 message() 与聚合按 self.aggr默认 add求和 out self.propagate(edge_index, xx, edge_weightedge_weight) return out self.biaspropagate是基类MessagePassing提供的通用骨架它根据edge_index把每个节点的特征复制到每条边的起点上调用你实现的message(x_j, ...)得到边消息再按目标节点聚合。于是设计一个新 GNN 层只剩下设计消息函数与聚合方式这正是 README 中 10 行代码实现EdgeConv的原因。EdgeConv把点云变成图点云本身没有边DGCNN 的做法是每一层前向时动态计算 k 近邻图。DynamicEdgeConvtorch_geometric/nn/conv/edge_conv.py继承自EdgeConv在forward中先调用knn_graph现场建图再做消息传递消息函数拼接[x_i, x_j - x_i]后经 MLP、max 聚合。在示例中它被组装成分类网络examples/dgcnn_classification.py 91–108 行class Net(torch.nn.Module): def __init__(self, out_channels, k20, aggrmax): super().__init__() # 动态边卷积每层前向时按 k20 的 kNN 图做消息传递 self.conv1 DynamicEdgeConv(MLP([2 * 3, 64, 64, 64]), k, aggr) self.conv2 DynamicEdgeConv(MLP([2 * 64, 128]), k, aggr) self.lin1 Linear(128 64, 1024) self.mlp MLP([1024, 512, 256, out_channels], dropout0.5, normNone) def forward(self, data): pos, batch data.pos, data.batch x1 self.conv1(pos, batch) # 输入即坐标边每层动态生成 x2 self.conv2(x1, batch) out self.lin1(torch.cat([x1, x2], dim1)) out global_max_pool(out, batch) # 按 batch 向量做全局池化 return F.log_softmax(self.mlp(out), dim1)注意batch张量的作用DataLoader把多个点云拼成一个大张量时data.batch记录每个点属于哪个图global_max_pool(out, batch)依此分组池化——这是 PyG 处理变长图批处理的标准方式。PointTransformerConv给注意力加上相对位置Point Transformer 层在注意力分数中引入了点间相对位移torch_geometric/nn/conv/point_transformer_conv.pypos_nn把pos_j - pos_i映射成位置编码attn_nn作用于特征差值两者相加后 softmax 得到注意力权重。完整组装方式见 examples/point_transformer_classification.py其骨干通过 FPS 降采样 kNN 聚合实现层次化结构同文件 54–84 行的TransitionDowndef forward(self, x, pos, batch): id_clusters fps(pos, ratioself.ratio, batchbatch) # 最远点采样选簇 # 为每个簇中心找 k 个最近点作为邻居 id_k_neighbor knn(pos, pos[id_clusters], kself.k, batch_xbatch, batch_ysub_batch) x self.mlp(x) # 把 k 个邻居特征 max 池化回簇中心完成降采样 x_out scatter(x[id_k_neighbor[1]], id_k_neighbor[0], dim0, dim_sizeid_clusters.size(0), reducemax) return x_out, pos[id_clusters], sub_batch该示例依赖pyg_lib扩展提供 GPU 加速的fps/knn算子见文件 21–22 行的检查逻辑这是核心装完即可用、扩展按需安装设计的一个具体体现。端到端实战用 DGCNN 分类 ModelNet10下面把前面机制串成一个可运行的完整案例。仓库提供的 examples/dgcnn_classification.py 支持modelnet10 / modelnet40 / medshapenet三种数据集这里以默认的 ModelNet1010 类 3D 模型训练集约 9840 个为例。数据准备两条变换即可完成预处理pre_transform在数据首次加载时标准化坐标transform每次读取时采样 1024 个点pre_transform, transform T.NormalizeScale(), T.SamplePoints(1024) train_dataset ModelNet(root, 10, True, transform, pre_transform) test_dataset ModelNet(root, 10, False, transform, pre_transform)运行python examples/dgcnn_classification.py --dataset modelnet10 --epochs 201预期输出每个 epoch 打印一行Epoch 001, Loss: x.xxxx, Test: x.xxxx对应示例 144–147 行的循环。ModelNet10 上训练约 200 轮后测试准确率通常达到 92% 以上若 GPU 显存不足可降低--batch_size。想换医疗点云数据只需--dataset medshapenet示例内置了按类别 7:3 分层划分训练/测试集的逻辑50–70 行。进阶玩法邻居采样训练与多 GPU 分布式邻居采样大图小批量训练。当整图超过显存时NeighborLoader对每个种子节点采样 1–2 层邻居子图构成 mini-batch这是 GraphSAGE 论文的工业级落地源码见 torch_geometric/loader/neighbor_loader.pyloader NeighborLoader( data, num_neighbors[25, 10], # 两层采样每层邻居数 25、10 batch_size1024, input_nodestrain_idx, # 只在训练节点上做监督 shuffleTrue, )适用场景Reddit、OGB 等百万级节点图此时注意 loss 只取out[:batch.batch_size]因为 batch 里还包含采样引入的邻居节点。多 GPU 分布式训练。仓库提供了单机多卡完整示例 examples/multi_gpu/distributed_sampling.pymp.spawn启动每个 GPU 一个进程每个进程用 NCCL 加入进程组后各训练自己的标签分片最后对准确率做dist.all_reduce平均def run(rank, world_size, dataset): dist.init_process_group(nccl, rankrank, world_sizeworld_size) data dataset[0] # 标签节点按 world_size 切片每张卡训练自己的分片 train_idx data.train_mask.nonzero(as_tupleFalse).view(-1) train_idx train_idx.split(ceil(train_idx.size(0) / world_size))[rank] ... model SAGE(dataset.num_features, 256, dataset.num_classes).to(rank) model DistributedDataParallel(model, device_ids[rank])若只是单机多卡且图能放进显存直接DistributedDataParallel包装模型配合DataLoader即可上面的邻居采样分布式路径则面向图大于一卡的场景。此外 examples/datapipe.py 演示了与torchdata.DataPipe的集成适合构建大规模数据流水线。小结PyG 的核心价值在于一个统一的数据与消息传递抽象Dataedge_index让任意图结构变成张量MessagePassing把设计 GNN 层压缩为设计消息与聚合两个函数再配合内置数据集、点云算子与分布式加载器从十几行的 Cora 实验到亿级图的分布式训练都走在同一条路径上。建议按本文顺序动手先跑通 GCN 快速示例建立直觉再对照源码读懂propagate然后复现 DGCNN 案例最后按需进入采样与分布式训练。延伸阅读官方文档与教程入口docs/source/自定义 GNN 层教程docs/source/tutorial/点云模型性能基准benchmark/points/【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考