一种面向基于采样的图神经网络训练的数据加载系统

1.本发明属于软件开发技术领域，尤其涉及一种面向基于采样的图神经网络训练的数据加载系统。


背景技术：

2.随着数据的不断增长，数据间的关系愈加复杂，图神经网络得到了广泛的关注。不同于传统的擅长处理欧式空间数据的深度神经网络和专门处理图数据的传统图计算，图神经网络专注于处理非欧空间数据，将神经网络中的自动微分和传统图计算中的消息传递机制相结合，在图数据处理上具有较好的效果，并成功的应用在实际生产环境中，例如，社交网络、交通预测、推荐系统等等。
3.为了提升图神经网络的学习能力和解决图神经网络在处理大图面临的内存限制问题，基于采样的图神经网络被广泛研究。主流的图采样方式可分为三类：基于节点采样、基于层采样和基于子图采样。然而这些采样方法仅从数学角度考虑，忽略了在系统中的运行效率。另一方面，现有的深度学习框架例如pytorch、tensorflow等缺乏高效训练图神经网络的能力，学术界和工业界均不断提出图神经网络框架。目前具有代表性的图神经网络框架：dgl（deepgraphlibrary）、pyg（pytorch
‑ꢀ
geometric）、aligraph。这些图神经网络框架实现了众多的采样算法和图神经网络模型，为部署设计图神经网络模型带来了便利。通常情况下，实际生活中的图数据往往具有较大的规模，图神经网络框架面临着巨大的计算和存储压力。基于采样的图神经网络很好的缓解了大规模图数据带来的存储压力。在网络训练之前，基于采样的图神经网络限制聚合邻居节点的范围以此减少图数据的规模，从而缓解了存储压力。考虑到图数据中普遍存在的邻居爆炸问题，图神经网络的网络结构相比于深度神经网络较小，同时随着gpu的算力不断增强，数据加载包括数据传输成为了图神经网络训练过程中的瓶颈。在分布式环境下，数据加载瓶颈愈加严重。
4.数据加载瓶颈问题增加了图神经网络端到端的训练时间，降低了整个图神经网络系统的效率，无法充分利用gpu资源。为了提高数据加载效率，现有的部分技术针对图神经网络中的数据加载进行了优化。例如，dgl支持在gpu上的采样以此加速采样，然而此种采样方法需要将整个图加载到gpu，扩展性受到gpu显存容量的限制；pytorch-direct设计了一种cpu和gpu统一的张量类型，利用零拷贝技术直接从cpu端获取数据，忽略了不同mini-batch之间存在冗余的数据传输；pagraph提前在gpu端缓存度高的节点减少cpu和gpu之间冗余的数据传输，而占用了部分珍贵的gpu显存资源；rbd充分利用不同mini-batch之间的共享数据，减少了冗余的数据传输，而对于未被共享的节点数据仍然需要从cpu端获取，同样存在一定的数据传输开销；torch-quiver采用缓存热数据的方式提升数据传输效率，占用了额外的gpu显存资源。
5.综上所述，数据加载成为提升图神经网络训练效率的最大限制性因素。现有的技术在采样、数据传输分别进行了优化，但是均存在其他方面的缺点，图神经网络训练中的数据加载需要进一步优化。


技术实现要素：

6.本发明结合图采样特性、数据传输特性以及图数据特征，实现了一个高效、易用、基于采样的图神经网络的数据加载系统，目的是解决图神经网络训练过程中面临的数据加载（图采样和数据传输）瓶颈问题。
7.为解决上述技术问题，实现本发明目的的具体技术方案是：一种面向基于采样的图神经网络训练的数据加载系统，该系统包括：邻居节点采样器和数据传输器，其中：所述数据传输器包括分类器、特征聚合器和数据管理器；所述邻居节点采样器与数据传输器中的分类器相连，利用采样算子采样得到当前mini-batch节点，同时采用一种csr数组存储策略减少采样过程中的数据传输时延；所述数据传输器中的分类器分别与邻居节点采样器和数据管理器相连，将邻居节点采样器采样得到的节点进行分类，分为共享节点和非共享节点；所述数据传输器对不同类节点进行不同的处理：对于共享节点，直接使用gpu中维持的特征数据，对于非共享节点，利用特征聚合器获取其特征数据；所述数据传输器中的数据管理器在gpu中维持一块可变大小的内存空间保存上一个mini-batch的特征数据并在每次训练迭代过程中就地更新该内存空间中的特征数据。
8.所述采样算子以每个节点的采样任务为基本单位开展并行采样，输出采样后得到的所有节点，形成一个mini-batch；所述csr数组存储策略，是根据csr数组大小将csr分别存储于shared memory、global memory和cpu memory中；当csr数组存放于cpu memory时，邻居节点采样器使用零拷贝进一步减少数据传输时延。
9.数据传输器中的分类器对邻居节点采样器输出的节点进行分类，采用倒排索引的方式将当前mini-batch中待获取特征数据的节点分为两类：共享节点和非共享节点；所述分类器对于上一个mini-batch和当前mini-batch的节点建立倒排索引，分别记录节点在两个mini-batch中的索引；根据索引个数将当前mini-batch节点分类为共享节点和非共享节点。
10.所述数据传输器的特征聚合器，首先利用零拷贝减少数据传输时延，其次将每个节点的特征聚合任务视为并行的基本单位展开并行加速；针对单个节点的特征聚合，内部以特征维度为单位继续开展并行。
11.所述数据传输器的数据管理器在gpu中维持一块可变大小的内存空间保存上一个mini-batch的特征数据，提供一种就地更新该特征数据的方式；若当前mini-batch的节点数多于上一个mini-batch的节点数，数据管理器对gpu中维持的内存空间扩容，反之则移除多余的空间；所述就地更新内存空间的方式在于将非共享节点的索引进行排序，再分别替换gpu中内存空间中的特征数据。
12.本发明的有益效果是，通过上述一种面向基于采样的图神经网络数据加载系统，减少了数据加载时延，提升了系统的整体性能。
附图说明
13.图1是本发明实施例的数据加载系统框图；图2是本发明实施例的数据加载系统使用流程图；图3是本发明实施例的数据加载系统中邻居节点采样器流程图；图4是本发明实施例的数据加载系统中数据传输器中处理不同类型数据流程图；图5是本发明实施例的数据加载系统中数据传输器中分类器流程图；图6是本发明实施例的数据加载系统中特征聚合流程图。
具体实施方式
14.下面结合实施例及附图，详细描述本发明的技术方案。
15.本发明的一种面向基于采样的图神经网络训练的数据加载系统，结合图采样、数据传输特性以及图数据特征对数据加载进行优化。
16.一种面向基于采样的图神经网络的数据加载系统，分别针对图采样和数据传输进行优化以此减少数据加载时延。所述系统包括邻居节点采样器和数据传输器。邻居节点采样器，利用gpu强大的并行计算能力，并提供一种csr存放策略提升采样速度；数据传输器包括分类器、特征聚合器和数据管理器，利用不同mini-batch之间存在重复采样节点的特性，分类器将当前mini-batch的节点分为共享节点和非共享节点；对于共享节点，无需重复传输其特征数据；对于非共享节点，利用特征聚合器从cpu获取特征数据；数据管理器在gpu中维持一块可变大小的内存空间保存上一个mini-batch的特征数据并采用就地更新的方式更新gpu中的特征数据；所述系统提供简单易用的操作接口，包括sampler（采样）、重载操作符获取节点特征。
17.所述系统中的邻居节点采样器将图采样过程从cpu迁移到gpu，在gpu端设计了一个高效的采样算子，利用gpu的高度并行性掩盖随机访存时延。采样算子以深度学习框架pytorch中的dataloader生成的种子节点为输入，输出对上述种子节点进行采样后的所有节点，形成一个mini-batch。采样算子以每个节点的采样任务为基本单位开展并行采样。邻居节点采样器提供一种csr数组存储策略：根据csr数组大小将csr分别存储于shared memory、global memory和cpu memory中。对于csr数组存放于cpu memory中，邻居节点采样器使用零拷贝技术进一步减少数据传输时延。
18.所述系统中的数据传输器利用分类器将mini-batch中待获取特征数据的节点分为两类：共享节点和非共享节点。对于共享节点，数据传输器直接使用保存在gpu中的特征数据，减少cpu和gpu之间冗余的数据传输；对于非共享节点，数据传输器设计了一个高效的gpu端的特征聚合器，利用gpu的并行性减少从cpu获取特征数据的时延。分类器采用倒排索引的方式将mini-batch中待获取特征数据的节点分为两类。分类器对于上一个mini-batch和当前的mini-batch节点建立倒排索引，分别记录节点在两个mini-batch中的索引。根据索引个数将mini-batch节点分类为共享节点和非共享节点。数据传输器利用数据管理器在gpu中维持一块可变大小的内存空间，提供了一种就地更新该内存空间的方式。当目前mini-batch的节点数多于上一个mini-batch的节点数，数据传输器对gpu中的内存空间扩容，反之则移除多余的空间。所述就地更新内存空间的方式在于将非共享节点的索引进行排序，再分别替换gpu中内存空间中的特征数据。
19.所述系统的使用步骤如下：步骤a-1：利用pytorch中的dataloader生成mini-batch的种子节点；步骤a-2：邻居节点采样器对种子节点进行高速采样，采样生成的全部邻居节点以及种子节点构成一个mini-batch；步骤a-3：数据传输器获取mini-batch中节点的特征数据；步骤a-4：开始训练。
实施例
20.本实施例的数据加载系统，其框图参阅图1。实施例系统中包含两个关键组件，邻居节点采样器和数据传输器。邻居节点采样器，采用csr存放策略以及为每个节点的采样任务分配一个warp的gpu线程的方式，提高图采样的效率。数据传输器针对不同mini-batch之间存在共享数据的特性，减少了cpu到gpu之间冗余的数据传输，并针对非共享节点，设计了一个高效的特征聚合器，进一步提高数据传输的吞吐。实施例系统提供简单易用的操作接口，包括sampler（采样）、重载操作符来获取节点特征，使得用户只需修改少量代码即可运行所述系统并获得系统性能提升。本发明实施例系统的流程见图2，其包括如下步骤：步骤101：利用深度学习框架pytorch中的dataloader生成种子节点；步骤102：邻居节点采样器接收种子节点，生成采样后的节点，发送给数据传输器中的分类器；步骤103：分类器依据收到的mini-batch节点与上一个mini-batch节点进行比较并分类，分类成共享节点和非共享节点；步骤104：数据管理器在gpu中维持一块可变大小的空间保存上一个mini-batch的特征数据；步骤105：数据传输器对于共享节点，直接使用保存在gpu中的特征数据；对于非共享节点，利用特征聚合器从cpu端获取数据；步骤106：数据管理器就地更新维持在gpu中的特征数据；步骤107：数据管理器将获取到的mini-batch的特征数据发送给模型，进行训练。
21.本发明实施例的邻居节点采样器，目的用于提高采样效率。邻居节点采样器将采样从cpu迁移到gpu，设计了一个高效的采样算子，利用gpu的高度并行性掩盖随机访存的时延。采样算子具体设计原则如下：对于一批种子节点而言，将每个节点的采样任务视为一个基本任务，每个warp负责一个基本任务，即一个节点的采样任务。在同一个warp中的每个线程分别负责采样种子节点的一个邻居。该设计使得同一个warp中的线程能够连续访存，从而减少gpu端发起的内存访问事物的数目，减少运行开销。由于图的结构信息通常以csr或者csc的方式存储以减少存储需求，图采样过程中每个线程需要访问csr数组获得邻居节点。邻居节点采样器提出一种基于csr的存放策略进一步减少采样时延。具体的，基于csr的存放策略如下：1、当size(csr)《= size(sharedmemory)：将csr数组存放于sharedmemory中。此种方法访存时延最低。
22.2、当size(csr) 》 size(shared memory) 同时 size(csr) 《= size(global memory)：将csr数组存放于global memory中。此时csr依旧保存在gpu端，时延比从cpu端拉
取数据更低。
23.3、当size(csr) 》 size(global memory)：此时gpu无法存放csr数组，采用零拷贝的方式从cpu获取图结构信息，相比于直接从cpu端获取数据，降低了数据传输的时延。
24.本实施例邻居节点采样器的流程图如图3所示，其包括以下步骤：步骤201：在系统开始运行之前，计算csr数组大小，依据csr数组存放策略决定csr的存放位置；步骤202：对于输入的一批种子节点，每个节点分配一个warp；步骤203：对于单个节点，warp中的每个线程负责采样其对应的种子节点的邻居节点；步骤204：采样完成，形成mini-batch节点编号。
25.实施例中的数据传输器，目的用于提高数据传输的吞吐。数据传输器包含三部分：分类器、特征聚合器和数据管理器。数据管理器在gpu维持一块空间来保存上一个mini-batch的特征数据并在每次训练迭代过程中就地更新特征数据；分类器将要获取特征数据的节点分为两类。第一类为不同mini-batch之间共享的节点，数据传输器直接复用保存在gpu中的此类节点的特征数据，减少cpu到gpu之间的数据传输量。第二类为不同mini-batch之间非共享的节点，数据传输器设计了一个高效的特征聚合器减少从cpu获取这类节点的特征数据的时延。图4展示了本实施例数据传输器处理不同类型数据的流程，具体包括如下步骤：步骤301：分类器对mini-batch节点进行分类；步骤302：对于共享节点，数据传输器直接使用依旧保存在gpu上的数据；步骤303：对于非共享节点，数据传输器利用特征聚合器从cpu端获取特征数据。
26.分类器采用倒排索引对两个不同mini-batch的节点进行分类，该倒排索引记录每个节点在两个相邻mini-batch中的索引。如果节点的倒排索引包含两个索引位置，则将其归类为共享节点；否则归类为非共享节点。在倒排索引的过程中，分类器会记录下共享节点的索引以及出现在两个mini-batch中非共享节点的索引，用于接下来的特征聚合过程。图5展示了本实施例数据传输器中分类器的实现流程图，具体包括如下步骤：步骤401：为上一个mini-batch中的节点和当前mini-batch的节点建立倒排索引；步骤402：根据倒排索引的数目对全部节点进行分类；具有两个索引的节点归类为共享节点，其余节点归类为非共享节点；步骤403：记录共享节点的索引和非共享节点在两个mini-batch中的索引。
27.特征聚合器实现了一个高效的特征聚合算子。该聚合算子的主要设计思想如下：首先利用零拷贝技术将特征数据保存在cpu端的锁页内存中，减少在cpu端的一次内存搬运。对于待获取特征数据的节点，将每个节点的特征聚合视为一个基本任务，每个warp负责一个基本任务，即一个节点的特征聚合任务。在同一个warp中的每个线程分别负责聚合一个维度的特征数据。该设计使得同一个warp中的线程能够连续访存，从而进一步的减少访存时延，减少运行开销。由于不同mini-batch可能会存在节点数目不一致的情况，数据管理器在gpu中维护一个可动态调整大小的存放节点特征数据的数组，采用就地更新的方式更新在gpu中的节点特征数组。为实现就地更新的方式，数据管理器首先对处于两个mini-batch中的非共享节点的索引进行排序，在mini-batch变小的情况下，将原先mini-batch中
后面多余的节点特征数据移除。排序后的索引能保证顺利有效的移除掉mini-batch中多余的节点特征数据。对于mini-batch变大的情况，将mini-batch数组进行扩容。上述就地更新方式使得最终获取到的mini-batch特征数据的顺序与最初的mini-batch节点顺序不一致，因此在获取特征数据后，需要将mini-batch的标签数据进行重排，使得标签数组和mini-batch特征数组相对应。本实施例特征聚合（特征聚合器和数据管理器配合使用）的流程图如图6所示，具体包括如下步骤：步骤501：将cpu端特征数据分配至锁页内存；步骤502：排序mini-batch中非共享节点的索引；步骤503：如果当前mini-batch的节点数大于上一个mini-batch的节点数，对gpu端的mini-batch数组扩容；步骤504：对于每个非共享节点，分别分配一个warp负责特征聚合；步骤505：在同一个warp中的每个线程负责一个维度的特征聚合；步骤506: 替换非共享节点特征数据；步骤507: 非共享节点特征替换结束后，如果当前mini-batch节点的数目小于上一个mini-batch的节点数目，则移除gpu中mini-batch数组中多余的节点数据。