一种基于全局索引表的快速特征采集方法

技术领域

本发明属于机器学习中神经网络算法应用领域，尤其涉及一种基于全局索引表的快速特征采集方法。

背景技术

图神经网络（GNN）是针对图这一数据结构的深度学习模型，它在多项基于图的机器学习任务中表现出色，并且具有较高的可解释性，近年来已成为一种广泛应用的图表示学习方法。GNN在广义上是指以图结构为输入的所有神经网络，其核心思想是学习一种映射，将图结构中的节点及其邻居节点的特征映射为该节点的低维隐特征表示。为了在GPU内存容量有限的场景下实现大规模图神经网络训练，基于采样的图神经网络训练被提出。图神经网络采样流程中的每一次迭代过程主要分为三个阶段，即采样、数据加载和GPU计算。

快速特征采集是指将缓存在GPU显存上的节点特征数据进行快速查找的过程（节点特征数据是图结构中的节点所对应的特征数据，图结构中每个节点都有其对应的特征数据），通过快速查找训练所需的节点特征数据来提高性能。图神经网络采样流程中的快速特征采集是数据加载过程能否快速进行的关键。

目前在学术界对于现有的图神经网络系统采用的训练算法有两大阵营，即全批训练和小批训练。全批训练的最主要代表有北京大学Ma等人提出的NeuGraph和斯坦福大学Jia等人提出的ROC。小批训练的最主要代表有IBM研究所Chen等人提出的FastGCN和南加州大学zeng等人提出的GraphSAINT。然而，对于像Facebook友谊图这样的非常大的图形，使用全批训练对GPU内存容量的要求将十分的高。但是GPU的内存容量增长却比较缓慢，难以满足大图在全批训练时对GPU内存容量的需求。因此小批量训练是如今对大图训练的一个非常重要的方法，小批量训练又称为基于采样的图神经网络训练。

在特征采集的过程中，当前这些方案基本采用了顺序查找和建立哈希表进行查找等技术。所述的顺序查找是按顶点数搜索特征，需要遍历所有缓存在GPU显存上的数据，因此顺序查找非常耗时。而所述的建立哈希表进行查找是根据关键码值进行访问，也就是通过把关键码值映射到表中一个位置来进行访问，但是哈希表是基于数组的，数组创建后难以扩展，当数组被填满后，性能会受到很大的影响。

他们的共同特点都是在GPU缓存中通过分配连续的内存块用以存储特征数据，并根据在GPU缓存中所创建的顺序表或者哈希表对特征数据进行查找，加快特征查找过程中所需特征数据的回复速度。所述的存储特征数据是一个将顶点的特征数据组织成几个大的[N,Ki]矩阵（其中N表示缓存的顶点数目，Ki是第i个特征名称字段下的特征维度），分配到GPU缓存中连续内存块的过程。它提供的信息能够指导编译器或者程序员去利用空闲的GPU缓存，通过提高特征数据的查找速度，从而实现图神经网络训练中更加高效的特征采集。

但是当前的图神经网络中的特征采集方案均无法带来快速的特征数据采集的效果，从而导致GPU处于等待的时间过长。其原因在于使用顺序查找特征数据时，需要遍历所有缓存在GPU上的特征数据，直到所需特征数据被采集，查找操作才算结束。使用哈希表查找特征数据时，当创建的数组被完全填满时，特征查找的性能将受到很大的影响。

发明内容

本发明的目的在于提供一种基于全局索引表的快速特征采集方法，旨在解决目前图神经网络中的特征采集机制在使用顺序查找特征数据时，需要遍历所有缓存在GPU上的特征数据，直到所需特征数据被采集，查找操作才算结束。使用哈希表查找特征数据时，当创建的数组被完全填满时，特征查找的性能将受到很大的影响。

本发明是这样实现的，一种基于全局索引表的快速特征采集方法包括如下步骤：

步骤一、剖析内存使用，确定可用于缓存分配的最大GPU内存量；

步骤二、在离线状态下，按出度数对顶点进行预排序。并在运行时选择高出度的顶点来填充GPU缓存；

步骤三、确定mini-batch输入的顶点数；

步骤四、确定全局索引表中局部顶点的位置；

步骤五、从GPU缓存数据中获取输入特征；

步骤六、当前批次特征采集结束。

进一步，所述的步骤一通过剖析第一个mini-batch训练之后，在训练期间检查空闲GPU内存的大小，来确定可用于缓存分配的最大GPU内存量。

进一步，步骤二所述的在离线状态下，按出度数对顶点进行预排序。并在运行时选择高出度的顶点来填充GPU缓存：

步骤一、依据neighbor sampling算法的独特访问模式，顶点的出度程度表示了它在整个epoch中被选择的概率，且顶点的出度越高，它就越有可能成为其他顶点的近邻。因此，在离线状态下，按出度数对顶点进行预排序；

步骤二、根据在离线状态下按出度数对顶点进行预排序步骤一中已经得到顶点出度数的排序结果，然后在运行时选择高出度数的顶点来填充GPU缓存。

进一步，所述的步骤三通过从CPU加载到GPU的mini-batch数据，确定其中输入的顶点数目。

进一步，步骤四所述的确定全局索引表中局部顶点的位置：

步骤一、根据全局索引表中key值对应的是全局顶点编号，通过顶点的编号找到该顶点在key值中所对应的位置；

步骤二、根据全局索引表中value值对应的是顶点特征在缓存数据中的位置，通过value值确定顶点的特征在缓存数据中的具体位置。

进一步，所述的步骤五通过从GPU缓存数据中获取输入特征，用以进行GPU计算操作；

进一步，所述的步骤六通过基于全局索引表的操作，用以快速进行特征数据采集过程。

本发明的有益效果如下：

（1）本发明从特征采集的源头上避免了顺序查找耗时长以及哈希查找性能受限的问题，利用全局索引表中的key值和value值达到快速查找特征数据，进而降低了特征采集过程中的时间消耗，提高了GPU计算资源的有效利用率；

（2）本发明通过创建全局索引表的方式，用以快速查找特征数据，相对于当前方案更加符合GNN训练过程中快速采集特征数据的体系结构特征，不仅能够有效地加速特征采集的执行过程，更能提高系统的整体运行效能；

（3）本发明可以面向不同的图神经网络训练过程中的特征查找，创建全局索引表用来快速查找顶点特征数据，适用于多种目标平台，具有通用性。

附图说明

图1是本发明实施例提供的基于全局索引表的快速特征采集方法的流程图。

图2是本发明实施例提供的基于采样的GNN模型训练流程（其中数据加载包括②－③）。

图3是本发明实施例提供的使用全局索引表查找特征数据的过程。

图4是本发明实施例提供的在GPU缓存中使用全局索引表进行快速的特征数据查找。

具体实施方式

为了使本发明的目的、技术方案及优点更加清楚明白，以下结合附图及实施例，对本发明进行进一步详细说明。应当理解，此处所描述的具体实施例仅仅用以解释本发明，并不用于限定本发明。

图1示出了本发明的基于全局索引表的快速特征采集方法的流程，如图所示，本发明是这样实现的，一种基于全局索引表的快速特征采集方法包括如下步骤：

S101：剖析内存使用，确定可用于缓存分配的最大GPU内存量；

S102：在离线状态下，按出度数对顶点进行预排序。并在运行时选择高出度的顶点来填充GPU缓存；

S103：确定mini-batch输入的顶点数；

S104：确定全局索引表中局部顶点的位置；

S105：从GPU缓存数据中获取输入特征；

S106：当前批次特征采集结束。

进一步，所述的步骤S101通过剖析第一个mini-batch训练之后，在训练期间检查空闲GPU内存的大小，来确定可用于缓存分配的最大GPU内存量。

进一步，步骤S102所述的在离线状态下，按出度数对顶点进行预排序。并在运行时选择高出度的顶点来填充GPU缓存：

步骤一、依据neighbor sampling算法的独特访问模式，顶点的出度程度表示了它在整个epoch中被选择的概率，且顶点的出度越高，它就越有可能成为其他顶点的近邻。因此，在离线状态下，按出度数对顶点进行预排序；

步骤二、根据在离线状态下按出度数对顶点进行预排序步骤一中已经得到顶点出度数的排序结果，然后在运行时选择高出度数的顶点来填充GPU缓存。

进一步，所述的步骤S103通过从CPU加载到GPU的mini-batch数据，确定其中输入的顶点数目。

进一步，步骤S104所述的确定全局索引表中局部顶点的位置：

步骤一、根据全局索引表中key值对应的是全局顶点编号，通过顶点的编号找到该顶点在key值中所对应的位置；

步骤二、根据全局索引表中value值对应的是顶点特征在缓存数据中的位置，通过value值确定顶点的特征在缓存数据中的具体位置。

进一步，所述的步骤S105通过从GPU缓存数据中获取输入特征，用以进行GPU计算操作；

进一步，所述的步骤S106通过基于全局索引表的操作，用以快速进行特征数据采集过程。

实施例：Facebook友谊图是由7亿个顶点和1000多亿条边组成。除了图结构数据之外，与顶点和边相关联的高维特征（通常范围从300到600）导致了更大的计算和存储复杂性。因此，在时间和资源的限制下，完全批量地制作一个完整的巨型图序列将不再有效，甚至不可行。此时一种典型的做法就是采样，它重复地从原始图形中采样子图作为小批量的输入，减少了单个小批量计算，同时仍然收敛到预期的精度。Facebook友谊图上实现的基于全局索引表的快速特征采集方法步骤如下。

步骤一、剖析内存使用，确定可用于缓存分配的最大GPU内存量；

为了避免对Facebook友谊图进行训练时与其他高优先级训练计算的资源竞争，需要估计可用于缓存分配的最大GPU内存量。为了实现这一点，本发明利用了在训练迭代中内存消耗相似的事实。这是因为基于采样的小批量训练使用几乎相同数量的数据样本作为输入，并且执行几乎相同数量的计算来训练每次迭代的共享GNN模型。因此，只需要在第一个mini-batch的训练之后，在训练期间检查空闲的GPU内存大小，就能确定可用于缓存分配的最大GPU内存量。

步骤二、在离线状态下，按出度数对顶点进行预排序。并在运行时选择高出度的顶点来填充GPU缓存：

对Facebook友谊图进行训练时，为了生成更好的模型，需要对训练的样本序列进行随机混洗，这使得在运行的过程中很难预测哪个顶点最有可能在下一个小批次中被访问。本发明根据neighbor sampling算法的独特访问模式，在离线状态下，按出度数对顶点进行预排序；

如图2所示，此流程为基于采样的GNN模型训练的一次运行过程，它主要包括采样、数据加载和GPU计算的三大步骤。在以上步骤中已经得到了顶点出度数的排序结果，因此，只需要在运行时将出度数高的顶点特征数据提前缓存到图2所示的GPU内存当中即可。

步骤三、确定mini-batch输入的顶点数；

如图2所示，mini-batch数据是通过采样操作从图存储服务器中所获得的，然后从CPU加载到GPU中。当mini-batch加载到GPU时，如图3中的步骤①所示，首先需要确定mini-batch所输入的顶点数。

步骤四、确定全局索引表中局部顶点的位置：

根据全局索引表中key值对应的是全局顶点编号，如图3中的步骤②所示，通过输入顶点所对应的编号找到该顶点在key值中所对应的位置；

根据全局索引表中value值对应的是顶点特征在缓存数据中的位置，如图3中的步骤②所示，通过输入顶点的value值确定顶点特征在缓存数据中的具体位置。

步骤五、从GPU缓存数据中获取输入特征；

在步骤四中，已经确定了输入顶点的key值和value值，如图3中的步骤③所示，可以通过对应的value值从缓存的特征数据中快速查找到所需要的特征数据。用以进行接下来的GPU计算操作。

步骤六、当前批次特征采集结束。

通过使用全局索引表查找特征数据，如图3所示，只需要确定好key值和value值就可以从缓存的特征数据中快速查询输入的顶点所对应的特征数据。从而达到快速完成特征数据查找的过程。

以上所提出的在GPU缓存中使用全局索引表进行快速的特征数据查找过程如图4所示。由于使用顺序查找，每次查找特征都需要遍历缓存的全部数据，太过于耗时。使用哈希表查找又会遇到当数组被填满后，性能会受到很大影响的尴尬场景。因此，提出使用全局索引表来进行特征数据查找，只需要确定输入顶点所对应的key值和value值，就可以对所需的顶点特征进行快速的查找。避免以上两种查询方式所带来的瓶颈问题。从而达到更加符合GNN训练过程中快速采集特征数据的体系结构特征，不仅能够有效地加速特征采集的执行过程，更能提高系统的整体运行效能。

上述虽然结合附图对本发明的具体实施方式进行了描述，但并非对本发明保护范围的限制，所属领域技术人员应该明白，在本发明的技术方案的基础上，本领域技术人员不需要付出创造性的劳动即可做出的各种修改或变形仍在本发明的保护范围之内。