图数据单机放不下怎么办用 PyTorch Geometric 的 NeighborLoader 把图神经网络训起来【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric节点数从 Cora 的 2,708 涨到 Reddit 的 23.2 万、papers100M 的 1.83 亿之后edge_index一个张量就能把显存撑爆全图一次 forward 的训练写法直接失效。PyTorch GeometricPyG是 PyTorch 生态里做图神经网络的库它的 NeighborLoader 用mini-batch 邻居采样把大图训练拆成一批批小子图再配合多机分区不需要把整张图装进内存也能训 GNN。为什么是它没有它你要自己维护邻接表、逐批扩展 K 跳邻居、做节点重映射、切出诱导子图——一条采样 pipeline 就是几周工作量有了它这条链路被压缩成NeighborLoader(data, num_neighbors...)一行采样、重映射、mini-batch 构造全部内建。跑通最小示例import torch import torch.nn.functional as F from torch_geometric.datasets import Planetoid import torch_geometric.transforms as T from torch_geometric.loader import NeighborLoader, NeighborSampler from torch_geometric.nn import GCNConv # Cora 只有 2708 个节点用来快速验证流程换成 Reddit 同样能跑 dataset Planetoid(rootdata/Planetoid, nameCora, transformT.NormalizeFeatures()) data dataset[0].to(cuda) # 两层 GCN结构上就是两次邻居聚合层数和后面的采样跳数要对应 class GCN(torch.nn.Module): def __init__(self, in_c, hid, out_c): super().__init__() self.conv1 GCNConv(in_c, hid) self.conv2 GCNConv(hid, out_c) def forward(self, x, edge_index, edge_attrNone): x self.conv1(x, edge_index, edge_attr).relu() return self.conv2(x, edge_index, edge_attr) model GCN(dataset.num_features, 32, dataset.num_classes).to(cuda) opt torch.optim.Adam(model.parameters(), lr0.01) # 核心每个 batch 只采出 128 个训练节点周围 2 跳的诱导子图而不是加载全图 train_loader NeighborLoader( data, num_neighbors[25, 25], batch_size128, input_nodesdata.train_mask) # 用张量 mask 指定本 epoch 的种子节点 # 测试时不切 mini-batch把 test 节点的子图一次性采出来再评估 test_loader NeighborSampler(data, num_neighbors[25, 25], input_nodesdata.test_mask) for epoch in range(30): for batch in train_loader: batch batch.to(cuda) out model(batch.x, batch.edge_index, batch.edge_attr) # 前 batch_size 行恰好是 128 个种子节点其余是采来的邻居loss 只算种子 loss F.cross_entropy(out[:batch.size], batch.y[:batch.size]) opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): test_batch next(iter(test_loader)).to(cuda) logits model(test_batch.x, test_batch.edge_index, test_batch.edge_attr) print((logits[:test_batch.batch_size].argmax(1) test_batch.y[:test_batch.batch_size]).float().mean())四个关键设计决策为什么只传train_mask就能训练input_nodes接受张量 mask、索引或整数列表loader 每轮从中抽 128 个节点当种子再沿num_neighbors[25,25]向外扩两跳每跳采 25 个邻居。图里绿色 Local 是本地节点、黄色 Remote 是采进来的邻居GCN 每次 forward 只见到这个小子图这就是mini-batch 训练在大图上的等价物。为什么 loss 要写out[:batch.size]采样结果的行顺序有保证——前batch_size行是种子节点后面是邻居。评估同理logits[:batch.batch_size]对齐种子标签不需要自己做索引映射。为什么测试用NeighborSampler而不是另一个 loader全图训练时examples/gcn.py可以整图一次推理但大图上对 test 节点直接切片会漏掉邻居特征。用NeighborSampler(data, input_nodesdata.test_mask)先把子图采出来推理和训练走同一条链路逻辑不会跑偏。边特征怎么进模型data.edge_attr归一化边权原样透传给GCNConv边上有阻抗、权重等多维特征时同样直接传入图里没有边特征就把参数留空。异构图多节点/边类型也原生支持num_neighbors传字典即可按边类型分别配采样数写法见 torch_geometric/loader/neighbor_loader.py 开头的文档字符串。采样跳数与网络层数对齐遇到什么场景模型 3 层、num_neighbors只给了 2 个数或者反向信息传不进去。怎么做num_neighbors的长度就是采样跳数应等于消息传递的层数另外directedTrue默认只保留原图方向的边两个都采样到的节点之间可能缺一条反向边层数多于跳数时把它改成directedFalse补上双向边# 3 层模型 → 3 跳采样deep 图建议补双向边 train_loader NeighborLoader(data, num_neighbors[25] * 3, directedFalse, batch_size128, input_nodesdata.train_mask)效果/注意directedFalse会多采一些边速度略降但信息更完整迭代慢的话可以外套一层 torch_geometric/loader/cache.py 的CacheLoader把采样子图缓存住省掉重复采样。百万节点图按 METIS 分区用 DistNeighborLoader 做多机采样遇到什么场景单张 GPU 塞不下全图特征矩阵本身就几十 GB单机采样方案失效。怎么做先用 METIS 按拓扑离线切分区——连接紧密的节点会尽量落在同一台机器采样时本地邻居绿色留在本机跨分区的远程邻居黄色走 RPC 自动拉取模型代码不用动换的只是入口和分区逻辑参数细节见 docs/source/tutorial/distributed_pyg.rst 和 examples/distributed/pyg/from torch_geometric.distributed import (partition_graph, DistNeighborLoader) # 离线METIS 按拓扑连通性把图切成 8 份跨机通信量最小化 partition_graph(data, rootdist_data, num_partitions8) # 训练每个 rank 只负责本分区的种子节点采样到远程邻居时经 RPC 自动获取 # torchrun --nproc_per_node8 train.py loader DistNeighborLoader( data, num_neighbors[25, 25], batch_size1024, input_nodespyg_dist.node_mask(train_mask, rankrank))效果/注意分区边界决定了跨机通信量METIS 的切法比随机切省得多torch_geometric.distributed在较新版本已标记 deprecated多机大规模场景官方更推荐看 examples/distributed/ 里 cuGraph 路线的示例选型前先确认版本。用 torch_geometric.compile 融合算子遇到什么场景模型能跑、采样也合理但每轮迭代还是慢。怎么做import torch_geometric model torch_geometric.compile(model) # 算子融合 内核优化改动一行效果/注意不同卷积在不同数据集上的实测加速在 1.4x2.2x 之间见 docs/source/advanced/compile.rst⚠️ 个别组合不升反降图中 gcnReddit 是 0.9x上线前用自己的负载跑一遍基线再决定。小图全量 forward 完全够用NeighborLoader反而引入采样开销稠密全连接图每点连所有点采样毫无收益这种数据直接nn.Linear更合理。下一步先跑 examples/multi_gpu/distributed_sampling.py 看单机多卡采样的完整工程写法再读 docs/source/tutorial/distributed_pyg.rst 把方案扩到多机。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考