基于高效采样算法的时序图神经网络系统 (二)

本期 IDP Inspiration 我们将继续为大家带来宦成颖博士针对其“基于高效采样算法的时序图神经网络系统”的研究成果的分享。

上一期文章中讲述了在时序图神经网络中的高效图采样算法，本期文章将重点讲述如何进行高效的时序图神经网络的训练。

现有的图数据规模极大，导致时序图神经网络的训练需要格外长的时间，因此使用多 GPU 进行训练变得成为尤为重要，如何有效地将多 GPU 用于时序图神经网络训练成为一个非常重要的研究议题。本文提供了两种方式来提高多 GPU 训练的性能，包括局部性感知的数据划分策略以及高效的任务调度策略。

首先，本文使用高效的数据划分策略将图数据以及点的特征向量进行有效的划分，来降低 GPU 之间额外的通信开销。点的特征向量占据很大的存储空间，因此我们将点的特征向量按照三个维度进行划分并且存储在不同的 GPU 中。

如图 1 所示，首先我们按照时序图的时间维度进行切分。在不同的时间范围中，将该时间范围内的点数据按照源点和汇点进行二维划分，其中每个 GPU 存储一段区间内源点的特征向量。对于边数据的存储，我们将边数据直接存储在 CPU memory 中。在进行训练时再将对应的边数据加载进相应的 GPU 中。

图 1. 局部性感知的数据划分策略

由于不同的数据存储在不同 GPU 中，如何进行高效的任务调度来降低通信开销将成为一个至关重要的问题。为解决该问题，我们提出了一种高效的任务调度策略：

首先我们的任务调度将存储在不同 GPU 上的数据进行调度，从而充分利用 GPU 之间 NVLink 高速通信的优势。

在进行训练时，由于每个 GPU 存储一段区间内的点特征向量，每个 GPU 需要传输当前边数据对应的汇点的特征向量。在 8 块 V100GPU 中，部分 GPU 之间存在两条 NVLink 高速链路，我们将充分利用这些 NVLink 链路来传输点的特征向量，从而提升传输效率。

以 8 块 GPU 为例，我们一共进行 4 轮的点数据交互。每一轮中，我们首先将编号为 0 和 1，2 和 3，4 和 5，6 和 7 的四对 GPU 进行数据交互。交换完成后，对每个 GPU 内的数据进行训练。接下来将编号 0 和 4，1 和 2，3 和 7，5 和 6 的四对 GPU 进行点数据交换。在进行 4 轮后，当前时间范围内的图数据将被得到训练。同时在数据交换中，我们充分使用了 GPU 之间的 2 条 NVLink 链路。

图 2. 数据与传输过程进行流水线化处理

其次，我们的高效任务调度将不同数据块上的训练以及 GPU 上的通信进行流水线化处理，以隐藏部分通信开销。如图 2 所示，我们将训练过程与数据传输过程进行流水线化处理。

具体来说，我们将 GPU 的显存划分为三部分：第一部分存储固定的数据（神经网络参数以及源点的特征向量），第二部分存储当前神经网络训练所需的数据（ 包括边数据和汇点的特征向量），第三部分用来进行数据传输及存储下一次训练所需的数据。在下一次训练时，我们交换第二部分和第三部分的数据，类似于 Ping-Pong 操作。在此基础上，我们可以将训练过程和数据传输重叠起来从而提高吞吐量、提高训练性能。

综上，针对时序图神经网络中出现的高额设备间通信开销问题，本文提出了局部性感知的数据划分策略和高效的任务调度策略，以缩减时序图神经网络训练时的通信开销。同时本文在时序图神经网络中引入了图采样，以进一步加速神经网络的训练并减少通信开销。基于上述通信缩减策略，本文提出了时序图神经网络系统 T-GCN。实验结果表明，T-GCN 实现了最高 7.9 倍的整体性能提升。在图采样性能上，本文提出的线段二分搜索采样算法能够实现最高 38.8 倍的采样性能提升。