[发明专利]一种图神经网络数据缓存方法、装置、设备及存储介质

说明书

本发明公开了一种图神经网络数据缓存方法、装置、设备及存储介质；在本方案中，在将图数据集中各顶点的特征数据缓存至GPU显存之前，为了避免GPU显存不能缓存全部特征数据，可先对各顶点进行聚类，将属于同一聚类的顶点分配至同一GPU，每个GPU只缓存每一部分顶点的l阶邻域(l是图神经网络层数)；并且，由于同一聚类内的邻域扩展会集中在该聚类内，因此这种根据顶点聚类信息划分到GPU并且缓存扩展邻域的方式，可以有效减少GPU缓存的数据量，保证全部顶点的特征数据均缓存至对应的GPU显存，避免GPU频繁访问主存储器，提高数据加载速度。

技术领域

本发明涉及数据缓存技术领域，更具体地说，涉及一种图神经网络数据缓存方法、装置、设备及存储介质。

背景技术

图数据作为一种非结构化的数据，已经广泛应用到推荐系统、社交网络、知识图谱等领域中，图神经网络已经成为处理图数据的有力工具。目前图神经网络模型训练有一个关键问题是从主存储器到GPU(graphics processing unit，图形处理器)显存的冗余数据加载。由于现实世界图数据中顶点的依赖性，图数据中不同训练顶点可能共享许多公共的邻居顶点。比如在引文网络图中，每个顶点是一篇论文，边表示论文之间的引用关系，一篇论文可以引用多篇文章，同一篇论文也可能被多个不同的论文引用，这种共同邻居顶点是因为图中的一个顶点很可能连接到多个训练顶点，使得扩展不同训练顶点的邻域时可能会重复采样部分顶点到样本批次中。这种不同训练批次之间的数据重叠模式加重了数据加载的负担。因此，对于每个训练迭代，加载的顶点数目会是原始图数据集顶点数目的几十倍。由于GPU只能在数据加载后开始训练，因此在图神经网络模型计算之前有很长的等待时间。

为了减少冗余的数据加载，现有方案利用有限的GPU显存缓存部分特征数据，训练期间直接从GPU取数据，而不用频繁的访问主存储器。但是，由于GPU显存不能缓存全部特征数据，因此训练中仍有一些缺失的数据需要从主存储器中获取，导致数据加载时间改进有限。因此，如何避免GPU频繁访问主存储器，提高数据加载速度是本领域技术人员需要解决的问题。

发明内容

本发明的目的在于提供一种图神经网络数据缓存方法、装置、设备及存储介质，以避免GPU频繁访问主存储器，提高数据加载速度。

为实现上述目的，本发明提供一种图神经网络数据缓存方法，包括：

通过图聚类算法对图数据集进行聚类，确定各顶点的聚类信息；

在每个聚类中确定训练顶点，并根据每个训练顶点的聚类信息向不同图形处理器分配对应的目标训练顶点；

对不同图形处理器的目标训练顶点进行邻域扩展，确定与每个图形处理器对应的邻域顶点；其中，邻域扩展的扩展层数与图神经网络层数一致；

从主存储器中获取与每个图形处理器对应的邻域顶点的特征数据，并将所述特征数据缓存至对应的图形处理器的显存。

其中，所述在每个聚类中确定训练顶点，包括：

获取顶点映射表；所述顶点映射表中记录了所述图数据集中聚类前各顶点与聚类后各顶点的对应关系；

根据所述图数据集中聚类前各顶点的顶点类型以及所述顶点映射表，在每个聚类中确定训练顶点。

其中，所述根据每个训练顶点的聚类信息向不同图形处理器分配对应的目标训练顶点，包括：

根据图形处理器的总数量以及每个训练顶点的聚类信息，将属于不同聚类的训练顶点平均分配至不同图形处理器，确定每个图形处理器对应的目标训练顶点；其中，属于同一聚类的训练顶点均分配至同一图形处理器。

其中，所述对不同图形处理器的目标训练顶点进行邻域扩展，确定与每个图形处理器对应的邻域顶点，包括：

确定图神经网络的层数；