PyG 远程后端Remote Backends完全指南借助 FeatureStore 与 GraphStore 将 GNN 扩展到单机内存之外【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric本指南深入解析 PyTorch GeometricPyG2.2 及更高版本提供的可扩展图机器学习基础设施通过FeatureStore特征存储与GraphStore图结构存储两个抽象接口将节点特征与图结构迁移到远程/外置存储配合采样器Sampler与数据加载器NodeLoader/LinkLoader让用户能够在远超单机内存容量的超大图上训练 GNN。读完本文你将掌握远程后端的背景动机、两大存储抽象的接口设计与 CRUD 用法、采样器与加载器的协同机制以及如何在当前仓库中动手实现自己的远程后端。引言为什么需要远程后端一个实例化的图神经网络其数据本质由两部分组成节点/边特征信息图中节点与边对应的稠密向量attribute图结构信息图中的节点以及连接这些节点的边。观察 GNN 的训练方式可以立刻得到一个结论当数据规模超出所选加速器accelerator如 GPU的可用内存时就必须放弃全图训练full-batch training改为在采样子图上训练mini-batch training。虽然这种方式给学习过程引入了随机性但它把加速器的内存需求降低到了采样子图的规模这正是经典 mini-batch GNN 训练范式的核心思想见图 1。图 1经典 mini-batch GNN 训练范式全图与特征整体存放在 CPU DRAM每轮迭代将采样子图及对应特征送入加速器。然而mini-batch 训练并非解决所有图学习扩展性问题的银弹。由于每一轮学习都要采样子图再送入加速器传统做法要求图和特征常驻在用户机器的 CPU DRAM 中。在大规模场景下这个要求会变得相当沉重采购一台拥有足够 CPU DRAM 来容纳整张图和特征的机器非常困难采用数据并行训练时需要把图和特征复制到每一个计算节点图和特征很容易超过单台机器的内存上限。因此要扩展到超出单机内存的极大图和特征就必须把这些数据结构移出内存out-of-core只让执行计算的那个节点处理采样子图。为了达成这一目标PyG 引入了两个核心抽象来分别存储特征信息与图结构见图 2特征存放在一个键值式的FeatureStore中它必须支持高效的随机访问图信息存放在一个GraphStore中它必须支持采样器对其进行高效采样。图 2图数据在远程存储与训练实例之间的布局。左侧分布式存储中Graph Store 负责节点与边结构Feature Store 负责节点/边张量右侧训练实例内采样子图与特征拼接后即为前向/反向传播所需的全部数据。在 PyG2.2 及更高版本中图数据被拆分为特征与结构两部分、这些信息被存放在可能远离实际训练节点的地方以及它们之间的交互——这一切对最终用户是完全透明的。只要FeatureStore与GraphStore被恰当地定义并牢记上文提到的性能要求剩下的工作全部由 PyG 处理。⚠️ 注意事项本文讨论的远程后端 API 仍处于演进之中PyG 团队会持续改进其易用性与通用性未来可能发生变化。目前FeatureStore与GraphStore仅支持异质图heterogeneous graphs且不支持边特征同质图与边特征的支持即将到来。这一点在仓库源码中亦有印证DataType.from_data见 sampler/base.py将(FeatureStore, GraphStore)元组识别为remote类型而 NodeLoader/LinkLoader 的返回类型为HeteroData。FeatureStore特征的键值抽象FeatureStore持有图中节点与边的特征。特征存储通常是图学习应用中最主要的存储瓶颈因为图布局信息即edge_index本身相对廉价每条边约 32 字节。PyG 为各种FeatureStore实现提供了一个公共接口使其能够接入核心学习 API。其抽象基类定义在 feature_store.py实现细节通过一套CRUD 风格的接口与 PyG 解耦。实现者主要需要覆写三个方法方法作用源码位置put_tensor(tensor, *args, **kwargs)向存储中同步写入一个特征张量返回是否成功feature_store.pyget_tensor(*args, **kwargs)从存储中同步读取一个特征张量feature_store.pyremove_tensor(*args, **kwargs)从存储中删除一个特征张量返回是否成功feature_store.py这三个方法都以TensorAttr见 feature_store.py为参数标识。TensorAttr包含三个字段其顺序即索引调用时属性必须给出的顺序group_name张量所属的分组名例如异质图中的节点类型attr_name张量在分组内的名字例如x或edge_attrindex张量行对应的节点索引可为torch.Tensor、numpy.ndarray、slice或单个整数。底层方法_put_tensor/_get_tensor/_remove_tensor是抽象方法由子类实现公开方法则在调用前通过TensorAttr.cast完成属性解析并要求属性必须被完整指定fully specified否则抛出ValueError。除 CRUD 外接口还提供了multi_get_tensor批量读取默认实现逐条调用get_tensor实现类可覆写以获得更高性能、get_tensor_size、get_all_tensor_attrs以及update_tensor默认先删后插等辅助方法。这一设计同时赋予用户pythonic 的接口来检查和修改FeatureStore中的元素。以下是文档给出的完整示例feature_store CustomFeatureStore() paper_features ... # [num_papers, num_paper_features] author_features ... # [num_authors, num_author_features] # 写入特征 feature_store[paper, x, None] paper_features feature_store[author, x, None] author_features # 访问特征 assert torch.equal(feature_store[paper, x], paper_features) assert torch.equal(feature_store[paper].x, paper_features) assert torch.equal(feature_store[author, x, 0:20], author_features[0:20])上述索引语法由FeatureStore.__getitem__/__setitem__见 feature_store.py与AttrView见 feature_store.py共同实现完全指定的键会直接产出张量部分指定的键会返回一个AttrView视图该视图可继续按属性名或索引取值也可通过调用store[group, attr]()强制触发 GET 操作。例如feature_store[paper].x正是先得到paper的视图、再以属性访问方式补全attr_name的链式写法。从设计意图看FeatureStore抽象做出如下关键假设见 feature_store.py特征可通过TensorAttr中指定的任意属性唯一标识实现者负责妥善处理这些假设——例如一个简单的内存实现可以把所有元数据值与特征索引拼接作为键值存储中的唯一键更复杂的实现可以基于元数据对特征做有趣的分区。常见的FeatureStore实现形态是键值存储key-value store例如memcached、LevelDB、RocksDB都是可行的性能选项。源码中还标注了未来的重要 TODO异步put与get功能。GraphStore面向高效采样的图结构抽象GraphStore持有定义节点间关系的边索引。其目标是以支持从根节点高效采样的方式存储图信息采样算法由开发者自行选择。与FeatureStore类似PyG 为各种GraphStore实现提供了接入核心学习 API 的公共接口但与FeatureStore不同的是GraphStore不需要对全部元素提供随机访问而需要定义一种能提供高效子图采样的表示。其抽象基类定义在 graph_store.py核心 CRUD 方法包括方法作用源码位置put_edge_index(edge_index, *args, **kwargs)以EdgeAttr指定的格式写入边索引graph_store.pyget_edge_index(*args, **kwargs)读取边索引找不到时抛出KeyErrorgraph_store.pyremove_edge_index(*args, **kwargs)删除边索引graph_store.py边索引通过EdgeAttr见 graph_store.py唯一标识其字段包括edge_type边类型在 PyG 中为源节点、关系类型、目标节点的三元组layout边表示格式取值为EdgeLayout枚举中的COO、CSC或CSR见 graph_store.pyis_sorted边索引是否按目标节点排序对 COO 有意义CSC 天然有序CSR 定义上即无序size该边类型的源/目标节点数量。GraphStore抽象的关键假设是边索引仅以 COO、CSC 或 CSR 格式表示且一旦存入即静态不变不支持动态修改这一点在源码 docstring 中有明确说明。接口还内置了布局转换能力coo()/csr()/csc()方法见 graph_store.py可在三种格式间转换内部通过_edge_to_layout/_edges_to_layout完成转换时可选择是否将结果回写存储storeTrue。测试用例 test_graph_store.py 覆盖了基本的读写与 COO↔CSR↔CSC 转换逻辑。接口用法示例如下graph_store CustomGraphStore() edge_index torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]]) # 写入边 graph_store[edge, coo] coo # 访问边 row, col graph_store[edge, coo] assert torch.equal(row, edge_index[0]) assert torch.equal(col, edge_index[1])常见的GraphStore实现是图数据库graph database例如Neo4j、TigerGraph、ArangoDB、Kùzu都是可行的性能选项。当前仓库提供了一个与Kùzu图数据库结合的示例见 examples/distributed/kuzuKùzu 是一个为查询速度与可扩展性而构建的进程内属性图数据库其 Python API 直接输出可接入 PyG 接口的FeatureStore与GraphStore从而允许直接在存储在 Kùzu 中的图上训练 GNN示例包含papers_100M场景约 1.11 亿节点、16 亿边的ogbn-papers100M数据在单机上配合远程后端使用。采样器与 GraphStore 紧密耦合图采样器graph sampler与给定的GraphStore紧密耦合它操作GraphStore从输入节点出发产出采样子图。不同采样算法实现在BaseSampler接口见 sampler/base.py背后默认情况下PyG 的默认内存采样器会把所有边索引从GraphStore拉取到训练节点内存中转换为压缩稀疏列CSC格式然后复用预构建的内存采样例程自定义采样器实现则可以选择覆写BaseSampler.sample_from_nodes见 sampler/base.py和/或BaseSampler.sample_from_edges见 sampler/base.py调用GraphStore的专有方法以获得效率提升例如直接在远程GraphStore上执行采样。sample_from_nodes接收一个NodeSamplerInput包含输入示例索引、种子节点索引、可选时间戳与输入节点类型返回SamplerOutput或HeteroSamplerOutputsample_from_edges则接收EdgeSamplerInput并支持可选的neg_sampling配置用于负采样。此外BaseSampler还提供edge_permutation属性见 sampler/base.py报告采样过程中对边顺序的置换供数据加载器还原原始边 ID。源码还给出重要提示采样器中存放的任何数据都会在数据加载 worker 间被复制每个 worker 持有采样器的独立实例因此建议限制采样器内保存的信息量。# CustomGraphSampler 知道如何在 CustomGraphStore 上采样 node_sampler CustomGraphSampler( graph_storegraph_store, num_neighbors[10, 20], ... )数据加载器NodeLoader 与 LinkLoaderPyG 并未为GraphStore定义必须实现的采样领域专用语言DSL相反采样器与GraphStore通过数据加载器紧密耦合在一起。PyG 开箱即用地提供了两种数据加载器NodeLoader见 node_loader.py从输入节点采样子图用于节点分类任务LinkLoader从一条边的任一侧采样子图用于链接预测任务。这两种加载器都以FeatureStore、GraphStore和一个图采样器为输入内部调用采样器的sample_from_nodes或sample_from_edges方法执行子图采样。NodeLoader的__init__签名见 node_loader.py支持dataData、HeteroData或(FeatureStore, GraphStore)元组、node_sampler、input_nodes异质图中需以(node_type, indices)元组形式传入、input_time、transform、transform_sampler_output、filter_per_worker自动推断过滤发生在 worker 子进程还是主进程、custom_cls远程后端场景下返回的自定义HeteroData类等参数其余**kwargs直接透传给torch.utils.data.DataLoader如batch_size、shuffle、drop_last、num_workers。核心用法如下# 不再传入 PyG data 对象而是传入 # FeatureStore 与 GraphStore 组成的元组作为输入数据 loader NodeLoader( data(feature_store, graph_store), node_samplernode_sampler, batch_size20, input_nodespaper, ) for batch in loader: pass在内部NodeLoader.collate_fn会调用node_sampler.sample_from_nodes完成采样随后filter_fn见 node_loader.py负责把采样结果与特征拼接对远程后端Tuple[FeatureStore, GraphStore]调用filter_custom_store/filter_custom_hetero_store从特征存储中取出采样节点对应的特征构造出Data或HeteroData对象并可附上n_id、e_id、batch、num_sampled_nodes、num_sampled_edges、input_id等元数据最终送到加速器。加载器还支持分布式场景当采样器为DistNeighborSampler时走专门的分布式过滤路径。整体架构组件如何协同工作从高层次看上述组件共同协作为 PyG 内扩展 GNN 训练提供支撑见图 3数据加载器准确地说是每个 worker借助一个BaseSampler向GraphStore发起采样请求收到响应后数据加载器随后向FeatureStore查询采样子图中节点与边对应的特征数据加载器从图结构与特征信息中构造最终的 mini-batch发送给加速器执行前向/反向传播循环往复直至收敛。图 3统一FeatureStore、GraphStore、图采样器与数据加载器的公共接口与数据流(1) 根节点 → (2) 采样节点 → (3) 采样节点 → (4) 采样节点特征。上述所有类都通过公共接口通信因此它们是可扩展、可泛化的并且易于与用户日常使用的 PyG 集成——DataType.from_data见 sampler/base.py将(FeatureStore, GraphStore)元组判定为remote数据类型的逻辑正是这一即插即用设计的直接体现。动手实践从零实现远程后端要开始使用这一扩展能力推荐按以下步骤进行阅读接口通读FeatureStorefeature_store.py与GraphStoregraph_store.py的抽象基类定义理解TensorAttr/EdgeAttr的字段语义与 CRUD 方法契约实现FeatureStore覆写_put_tensor/_get_tensor/_remove_tensor保证对随机访问的高效支持键值存储是首选形态实现GraphStore覆写_put_edge_index/_get_edge_index/_remove_edge_index保证对子图采样的高效支持图数据库是首选形态实现采样器根据采样算法覆写BaseSampler.sample_from_nodes与/或sample_from_edges尽量把采样逻辑下沉到远程GraphStore侧以省去整图搬运接入加载器将三者作为参数传给NodeLoader或LinkLoader其余 PyG 功能将像纯内存应用一样无缝工作。一旦FeatureStore、GraphStore和BaseSampler实现正确只需把它们作为参数传递给NodeLoader或LinkLoaderPyG 的其余部分便会无缝运行与任何纯内存应用别无二致。可参考的仓库样例包括examples/distributed/kuzu/README.mdKùzu 远程后端示例说明覆盖 PubMed 与papers_100M约 1.11 亿节点 / 16 亿边两种规模examples/distributed/kuzu/papers_100Mogbn-papers100M大规模图在单机上的远程后端训练代码torch_geometric/data/feature_store.py 与 torch_geometric/data/graph_store.py两个抽象基类的完整定义与接口契约torch_geometric/sampler/base.pyBaseSampler、NodeSamplerInput、EdgeSamplerInput与输出类型定义torch_geometric/loader/node_loader.py 与 torch_geometric/loader/link_loader.py两种开箱即用的远程后端数据加载器test/data/test_graph_store.pyGraphStore读写与格式转换的测试用例可作为实现正确性的验证参考。结语PyG 的远程后端通过FeatureStore与GraphStore两个简洁、易用且可扩展的抽象将特征与结构的存储彻底解耦并以数据加载器为纽带与采样器紧密协作从而把可扩展 GNN 训练的复杂度从用户侧完全抽离。值得注意的是该特性仍处于密集开发阶段目前仅支持异质图、暂不支持边特征API 细节未来仍可能调整。如果你在使用过程中有任何问题、意见或顾虑可以前往 PyG 的 GitHub Discussions 或 Slack 与 PyG 核心团队交流共同推动这一方向的演进。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 PyG 远程后端（Remote Backends）完全指南：借助 FeatureStore 与 GraphStore 将 GNN 扩展到单机内存之外
PyG 远程后端Remote Backends完全指南借助 FeatureStore 与 GraphStore 将 GNN 扩展到单机内存之外【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric本指南深入解析 PyTorch GeometricPyG2.2 及更高版本提供的可扩展图机器学习基础设施通过FeatureStore特征存储与GraphStore图结构存储两个抽象接口将节点特征与图结构迁移到远程/外置存储配合采样器Sampler与数据加载器NodeLoader/LinkLoader让用户能够在远超单机内存容量的超大图上训练 GNN。读完本文你将掌握远程后端的背景动机、两大存储抽象的接口设计与 CRUD 用法、采样器与加载器的协同机制以及如何在当前仓库中动手实现自己的远程后端。引言为什么需要远程后端一个实例化的图神经网络其数据本质由两部分组成节点/边特征信息图中节点与边对应的稠密向量attribute图结构信息图中的节点以及连接这些节点的边。观察 GNN 的训练方式可以立刻得到一个结论当数据规模超出所选加速器accelerator如 GPU的可用内存时就必须放弃全图训练full-batch training改为在采样子图上训练mini-batch training。虽然这种方式给学习过程引入了随机性但它把加速器的内存需求降低到了采样子图的规模这正是经典 mini-batch GNN 训练范式的核心思想见图 1。图 1经典 mini-batch GNN 训练范式全图与特征整体存放在 CPU DRAM每轮迭代将采样子图及对应特征送入加速器。然而mini-batch 训练并非解决所有图学习扩展性问题的银弹。由于每一轮学习都要采样子图再送入加速器传统做法要求图和特征常驻在用户机器的 CPU DRAM 中。在大规模场景下这个要求会变得相当沉重采购一台拥有足够 CPU DRAM 来容纳整张图和特征的机器非常困难采用数据并行训练时需要把图和特征复制到每一个计算节点图和特征很容易超过单台机器的内存上限。因此要扩展到超出单机内存的极大图和特征就必须把这些数据结构移出内存out-of-core只让执行计算的那个节点处理采样子图。为了达成这一目标PyG 引入了两个核心抽象来分别存储特征信息与图结构见图 2特征存放在一个键值式的FeatureStore中它必须支持高效的随机访问图信息存放在一个GraphStore中它必须支持采样器对其进行高效采样。图 2图数据在远程存储与训练实例之间的布局。左侧分布式存储中Graph Store 负责节点与边结构Feature Store 负责节点/边张量右侧训练实例内采样子图与特征拼接后即为前向/反向传播所需的全部数据。在 PyG2.2 及更高版本中图数据被拆分为特征与结构两部分、这些信息被存放在可能远离实际训练节点的地方以及它们之间的交互——这一切对最终用户是完全透明的。只要FeatureStore与GraphStore被恰当地定义并牢记上文提到的性能要求剩下的工作全部由 PyG 处理。⚠️ 注意事项本文讨论的远程后端 API 仍处于演进之中PyG 团队会持续改进其易用性与通用性未来可能发生变化。目前FeatureStore与GraphStore仅支持异质图heterogeneous graphs且不支持边特征同质图与边特征的支持即将到来。这一点在仓库源码中亦有印证DataType.from_data见 sampler/base.py将(FeatureStore, GraphStore)元组识别为remote类型而 NodeLoader/LinkLoader 的返回类型为HeteroData。FeatureStore特征的键值抽象FeatureStore持有图中节点与边的特征。特征存储通常是图学习应用中最主要的存储瓶颈因为图布局信息即edge_index本身相对廉价每条边约 32 字节。PyG 为各种FeatureStore实现提供了一个公共接口使其能够接入核心学习 API。其抽象基类定义在 feature_store.py实现细节通过一套CRUD 风格的接口与 PyG 解耦。实现者主要需要覆写三个方法方法作用源码位置put_tensor(tensor, *args, **kwargs)向存储中同步写入一个特征张量返回是否成功feature_store.pyget_tensor(*args, **kwargs)从存储中同步读取一个特征张量feature_store.pyremove_tensor(*args, **kwargs)从存储中删除一个特征张量返回是否成功feature_store.py这三个方法都以TensorAttr见 feature_store.py为参数标识。TensorAttr包含三个字段其顺序即索引调用时属性必须给出的顺序group_name张量所属的分组名例如异质图中的节点类型attr_name张量在分组内的名字例如x或edge_attrindex张量行对应的节点索引可为torch.Tensor、numpy.ndarray、slice或单个整数。底层方法_put_tensor/_get_tensor/_remove_tensor是抽象方法由子类实现公开方法则在调用前通过TensorAttr.cast完成属性解析并要求属性必须被完整指定fully specified否则抛出ValueError。除 CRUD 外接口还提供了multi_get_tensor批量读取默认实现逐条调用get_tensor实现类可覆写以获得更高性能、get_tensor_size、get_all_tensor_attrs以及update_tensor默认先删后插等辅助方法。这一设计同时赋予用户pythonic 的接口来检查和修改FeatureStore中的元素。以下是文档给出的完整示例feature_store CustomFeatureStore() paper_features ... # [num_papers, num_paper_features] author_features ... # [num_authors, num_author_features] # 写入特征 feature_store[paper, x, None] paper_features feature_store[author, x, None] author_features # 访问特征 assert torch.equal(feature_store[paper, x], paper_features) assert torch.equal(feature_store[paper].x, paper_features) assert torch.equal(feature_store[author, x, 0:20], author_features[0:20])上述索引语法由FeatureStore.__getitem__/__setitem__见 feature_store.py与AttrView见 feature_store.py共同实现完全指定的键会直接产出张量部分指定的键会返回一个AttrView视图该视图可继续按属性名或索引取值也可通过调用store[group, attr]()强制触发 GET 操作。例如feature_store[paper].x正是先得到paper的视图、再以属性访问方式补全attr_name的链式写法。从设计意图看FeatureStore抽象做出如下关键假设见 feature_store.py特征可通过TensorAttr中指定的任意属性唯一标识实现者负责妥善处理这些假设——例如一个简单的内存实现可以把所有元数据值与特征索引拼接作为键值存储中的唯一键更复杂的实现可以基于元数据对特征做有趣的分区。常见的FeatureStore实现形态是键值存储key-value store例如memcached、LevelDB、RocksDB都是可行的性能选项。源码中还标注了未来的重要 TODO异步put与get功能。GraphStore面向高效采样的图结构抽象GraphStore持有定义节点间关系的边索引。其目标是以支持从根节点高效采样的方式存储图信息采样算法由开发者自行选择。与FeatureStore类似PyG 为各种GraphStore实现提供了接入核心学习 API 的公共接口但与FeatureStore不同的是GraphStore不需要对全部元素提供随机访问而需要定义一种能提供高效子图采样的表示。其抽象基类定义在 graph_store.py核心 CRUD 方法包括方法作用源码位置put_edge_index(edge_index, *args, **kwargs)以EdgeAttr指定的格式写入边索引graph_store.pyget_edge_index(*args, **kwargs)读取边索引找不到时抛出KeyErrorgraph_store.pyremove_edge_index(*args, **kwargs)删除边索引graph_store.py边索引通过EdgeAttr见 graph_store.py唯一标识其字段包括edge_type边类型在 PyG 中为源节点、关系类型、目标节点的三元组layout边表示格式取值为EdgeLayout枚举中的COO、CSC或CSR见 graph_store.pyis_sorted边索引是否按目标节点排序对 COO 有意义CSC 天然有序CSR 定义上即无序size该边类型的源/目标节点数量。GraphStore抽象的关键假设是边索引仅以 COO、CSC 或 CSR 格式表示且一旦存入即静态不变不支持动态修改这一点在源码 docstring 中有明确说明。接口还内置了布局转换能力coo()/csr()/csc()方法见 graph_store.py可在三种格式间转换内部通过_edge_to_layout/_edges_to_layout完成转换时可选择是否将结果回写存储storeTrue。测试用例 test_graph_store.py 覆盖了基本的读写与 COO↔CSR↔CSC 转换逻辑。接口用法示例如下graph_store CustomGraphStore() edge_index torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]]) # 写入边 graph_store[edge, coo] coo # 访问边 row, col graph_store[edge, coo] assert torch.equal(row, edge_index[0]) assert torch.equal(col, edge_index[1])常见的GraphStore实现是图数据库graph database例如Neo4j、TigerGraph、ArangoDB、Kùzu都是可行的性能选项。当前仓库提供了一个与Kùzu图数据库结合的示例见 examples/distributed/kuzuKùzu 是一个为查询速度与可扩展性而构建的进程内属性图数据库其 Python API 直接输出可接入 PyG 接口的FeatureStore与GraphStore从而允许直接在存储在 Kùzu 中的图上训练 GNN示例包含papers_100M场景约 1.11 亿节点、16 亿边的ogbn-papers100M数据在单机上配合远程后端使用。采样器与 GraphStore 紧密耦合图采样器graph sampler与给定的GraphStore紧密耦合它操作GraphStore从输入节点出发产出采样子图。不同采样算法实现在BaseSampler接口见 sampler/base.py背后默认情况下PyG 的默认内存采样器会把所有边索引从GraphStore拉取到训练节点内存中转换为压缩稀疏列CSC格式然后复用预构建的内存采样例程自定义采样器实现则可以选择覆写BaseSampler.sample_from_nodes见 sampler/base.py和/或BaseSampler.sample_from_edges见 sampler/base.py调用GraphStore的专有方法以获得效率提升例如直接在远程GraphStore上执行采样。sample_from_nodes接收一个NodeSamplerInput包含输入示例索引、种子节点索引、可选时间戳与输入节点类型返回SamplerOutput或HeteroSamplerOutputsample_from_edges则接收EdgeSamplerInput并支持可选的neg_sampling配置用于负采样。此外BaseSampler还提供edge_permutation属性见 sampler/base.py报告采样过程中对边顺序的置换供数据加载器还原原始边 ID。源码还给出重要提示采样器中存放的任何数据都会在数据加载 worker 间被复制每个 worker 持有采样器的独立实例因此建议限制采样器内保存的信息量。# CustomGraphSampler 知道如何在 CustomGraphStore 上采样 node_sampler CustomGraphSampler( graph_storegraph_store, num_neighbors[10, 20], ... )数据加载器NodeLoader 与 LinkLoaderPyG 并未为GraphStore定义必须实现的采样领域专用语言DSL相反采样器与GraphStore通过数据加载器紧密耦合在一起。PyG 开箱即用地提供了两种数据加载器NodeLoader见 node_loader.py从输入节点采样子图用于节点分类任务LinkLoader从一条边的任一侧采样子图用于链接预测任务。这两种加载器都以FeatureStore、GraphStore和一个图采样器为输入内部调用采样器的sample_from_nodes或sample_from_edges方法执行子图采样。NodeLoader的__init__签名见 node_loader.py支持dataData、HeteroData或(FeatureStore, GraphStore)元组、node_sampler、input_nodes异质图中需以(node_type, indices)元组形式传入、input_time、transform、transform_sampler_output、filter_per_worker自动推断过滤发生在 worker 子进程还是主进程、custom_cls远程后端场景下返回的自定义HeteroData类等参数其余**kwargs直接透传给torch.utils.data.DataLoader如batch_size、shuffle、drop_last、num_workers。核心用法如下# 不再传入 PyG data 对象而是传入 # FeatureStore 与 GraphStore 组成的元组作为输入数据 loader NodeLoader( data(feature_store, graph_store), node_samplernode_sampler, batch_size20, input_nodespaper, ) for batch in loader: pass在内部NodeLoader.collate_fn会调用node_sampler.sample_from_nodes完成采样随后filter_fn见 node_loader.py负责把采样结果与特征拼接对远程后端Tuple[FeatureStore, GraphStore]调用filter_custom_store/filter_custom_hetero_store从特征存储中取出采样节点对应的特征构造出Data或HeteroData对象并可附上n_id、e_id、batch、num_sampled_nodes、num_sampled_edges、input_id等元数据最终送到加速器。加载器还支持分布式场景当采样器为DistNeighborSampler时走专门的分布式过滤路径。整体架构组件如何协同工作从高层次看上述组件共同协作为 PyG 内扩展 GNN 训练提供支撑见图 3数据加载器准确地说是每个 worker借助一个BaseSampler向GraphStore发起采样请求收到响应后数据加载器随后向FeatureStore查询采样子图中节点与边对应的特征数据加载器从图结构与特征信息中构造最终的 mini-batch发送给加速器执行前向/反向传播循环往复直至收敛。图 3统一FeatureStore、GraphStore、图采样器与数据加载器的公共接口与数据流(1) 根节点 → (2) 采样节点 → (3) 采样节点 → (4) 采样节点特征。上述所有类都通过公共接口通信因此它们是可扩展、可泛化的并且易于与用户日常使用的 PyG 集成——DataType.from_data见 sampler/base.py将(FeatureStore, GraphStore)元组判定为remote数据类型的逻辑正是这一即插即用设计的直接体现。动手实践从零实现远程后端要开始使用这一扩展能力推荐按以下步骤进行阅读接口通读FeatureStorefeature_store.py与GraphStoregraph_store.py的抽象基类定义理解TensorAttr/EdgeAttr的字段语义与 CRUD 方法契约实现FeatureStore覆写_put_tensor/_get_tensor/_remove_tensor保证对随机访问的高效支持键值存储是首选形态实现GraphStore覆写_put_edge_index/_get_edge_index/_remove_edge_index保证对子图采样的高效支持图数据库是首选形态实现采样器根据采样算法覆写BaseSampler.sample_from_nodes与/或sample_from_edges尽量把采样逻辑下沉到远程GraphStore侧以省去整图搬运接入加载器将三者作为参数传给NodeLoader或LinkLoader其余 PyG 功能将像纯内存应用一样无缝工作。一旦FeatureStore、GraphStore和BaseSampler实现正确只需把它们作为参数传递给NodeLoader或LinkLoaderPyG 的其余部分便会无缝运行与任何纯内存应用别无二致。可参考的仓库样例包括examples/distributed/kuzu/README.mdKùzu 远程后端示例说明覆盖 PubMed 与papers_100M约 1.11 亿节点 / 16 亿边两种规模examples/distributed/kuzu/papers_100Mogbn-papers100M大规模图在单机上的远程后端训练代码torch_geometric/data/feature_store.py 与 torch_geometric/data/graph_store.py两个抽象基类的完整定义与接口契约torch_geometric/sampler/base.pyBaseSampler、NodeSamplerInput、EdgeSamplerInput与输出类型定义torch_geometric/loader/node_loader.py 与 torch_geometric/loader/link_loader.py两种开箱即用的远程后端数据加载器test/data/test_graph_store.pyGraphStore读写与格式转换的测试用例可作为实现正确性的验证参考。结语PyG 的远程后端通过FeatureStore与GraphStore两个简洁、易用且可扩展的抽象将特征与结构的存储彻底解耦并以数据加载器为纽带与采样器紧密协作从而把可扩展 GNN 训练的复杂度从用户侧完全抽离。值得注意的是该特性仍处于密集开发阶段目前仅支持异质图、暂不支持边特征API 细节未来仍可能调整。如果你在使用过程中有任何问题、意见或顾虑可以前往 PyG 的 GitHub Discussions 或 Slack 与 PyG 核心团队交流共同推动这一方向的演进。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 RELATED READING
延伸阅读
更多一线实战笔记与深度复盘，助您持续精进
