本期 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 倍 的采样性能晋升。