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

技术特征：
1.一种面向基于采样的图神经网络训练的数据加载系统，其特征在于，包括：邻居节点采样器和数据传输器，其中：所述数据传输器包括分类器、特征聚合器和数据管理器；所述邻居节点采样器与数据传输器中的分类器相连，利用采样算子采样得到当前mini-batch节点，同时采用一种csr数组存储策略减少采样过程中的数据传输时延；所述数据传输器中的分类器分别与邻居节点采样器和数据管理器相连，将邻居节点采样器采样得到的节点进行分类，分为共享节点和非共享节点；所述数据传输器对不同类节点进行不同的处理：对于共享节点，直接使用gpu中维持的特征数据，对于非共享节点，利用特征聚合器获取其特征数据；所述数据传输器中的数据管理器在gpu中维持一块可变大小的内存空间保存上一个mini-batch的特征数据并在每次训练迭代过程中就地更新该内存空间中的特征数据。2.如权利要求1所述的一种面向基于采样的图神经网络训练的数据加载系统，其特征在于，所述采样算子以每个节点的采样任务为基本单位开展并行采样，输出采样后得到的所有节点，形成一个mini-batch；所述csr数组存储策略，是根据csr数组大小将csr分别存储于share dmemory、global memory和cpu memory中；当csr数组存放于cpu memory时，邻居节点采样器使用零拷贝进一步减少数据传输时延。3.如权利要求1所述的一种面向基于采样的图神经网络训练的数据加载系统，其特征在于，数据传输器中的分类器对邻居节点采样器输出的节点进行分类，采用倒排索引的方式将当前mini-batch中待获取特征数据的节点分为两类：共享节点和非共享节点；所述分类器对于上一个mini-batch和当前mini-batch的节点建立倒排索引，分别记录节点在两个mini-batch中的索引；根据索引个数将当前mini-batch节点分类为共享节点和非共享节点。4.如权利要求1所述的一种面向基于采样的图神经网络训练的数据加载系统，其特征在于，所述数据传输器的特征聚合器，首先利用零拷贝减少数据传输时延，其次将每个节点的特征聚合任务视为并行的基本单位展开并行加速；针对单个节点的特征聚合，内部以特征维度为单位继续开展并行。5.如权利要求1所述的一种面向基于采样的图神经网络训练的数据加载系统，其特征在于，所述数据传输器的数据管理器在gpu中维持一块可变大小的内存空间保存上一个mini-batch的特征数据，提供一种就地更新该特征数据的方式；若当前mini-batch的节点数多于上一个mini-batch的节点数，数据管理器对gpu中维持的内存空间扩容，反之则移除多余的空间；所述就地更新内存空间的方式在于将非共享节点的索引进行排序，再分别替换gpu中内存空间中的特征数据。

技术总结
本发明公开了一种面向基于采样的图神经网络训练的数据加载系统，该数据加载系统包括：邻居节点采样器和数据传输器；邻居节点采样器，以深度学习框架Pytorch的Dataloader的输出为输入，利用采样算子采样邻居节点；数据传输器包含分类器、特征聚合器和数据管理器三部分；分类器对邻居节点采样器输出的节点进行分类，分为共享节点和非共享节点；数据管理器在GPU中维持上一个mini-batch的特征数据，同时提供一种就地更新GPU中特征数据的方式；特征聚合器设计一个高性能的特征聚合算子从CPU获取非共享节点的特征；对于共享节点，数据传输器直接使用维持在GPU中的特征数据。本发明提高了采样效率，减少了不必要的数据传输，提高了数据传输吞吐。高了数据传输吞吐。高了数据传输吞吐。


技术研发人员：熊颖彤 翁楚良
受保护的技术使用者：华东师范大学
技术研发日：2022.06.08
技术公布日：2022/8/12