如何进行高效的时序图神经网络的训练

现有的图数据规模极大，导致时序图神经网络的训练需要格外长的时间，因此使用多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倍的采样性能提升。

原作者：追求卓越的 Baihai IDP