图（Graph）数据在现实世界中非常常见，例如社交网络、交通网络、物理系统等等，近几年图神经网络的发展将图数据的分析与深度神经网络结合，在越来越多的领域发挥出重要的作用，例如电商推荐、生物化学结构分析、反恐反诈风险控制等等。数据规模也呈现越来越大之势，动辄上千万节点规模的图，让很多图神经网络的训练需要探寻并行计算的方式以加速。
然而，图数据与视觉、自然语言等领域的数据不同，没有划分好的训练样本，很多情况下训练样本是从一个巨大的图中去随机采样，特征也需要在采样的过程中遍历邻接节点，这种天然特性使得对大规模图的并行训练并不容易。
幻方AI最近尝试对图神经网络的并行训练进行优化，以 OGB 图数据集作为训练样本，在幻方萤火集群上复现 DeepGCNs 的实验。基于幻方自研的3FS、hfreduce等工具，采用灵活的分布式训练方法，进行不同任务场景下的 GNN 训练提速，取得预期效果。本期文章将为大家详细描述。
论文标题：DeepGCNs: Can GCNs Go as Deep as CNNs?
原文地址：https://www.deepgcns.org/
模型仓库：https://github.com/HFAiLab/Distributed-DeepGCNs
数据集仓库: https://doc.hfai.high-flyer.cn/api/datasets.html#hfai.datasets.OGB
模型介绍
DeepGCNs 集成了近两年一系列前沿的 GNN 工作，其旨在构建如 CNN 一样的深层图卷积网络，解决优化深层结构中的一系列问题，以达到更优的表现效果。具体包括：
- 尝试构建深层的图卷积网络（GCN），探索深层 GCN 结构导致的信息传递快，容易梯度消息、过渡平滑和过渡拟合的问题；
- 设计提出残差、密集连接和空洞卷积等方案，使得深层图卷积网络表示更加稳定；
- 设计提出使用新的聚合、归一化函数，以统一信息的传递；
- 研究逆连接、组卷积、权重绑定和平衡模型，以提高深层图卷积模型的内存和参数效率。
DeepGCNs 的结构如下图所示：
在 OGB 公开图数据集的很多任务测试中，DeepGCNs 都取得了优异的效果。
可以看到，大规模图与深度图卷积模型的结合在越来越多的场景中发挥出惊人的效果，在 OGB 的公开结果测试排名中可以看到，对显卡算力的需求越来越大。如何高效的利用计算资源，尽可能加速图卷积网络的训练，是之后走向落地商用的必研究课题。
图的分布式训练
那么，如何对图数据进行分布式训练呢？在这之前，我们先看下图领域下的三大任务：
- 点属性预测（Node Property Prediction）
- 边属性预测（Link Property Prediction），也叫作链接预测
- 图属性预测（Graph Property Prediction）
这里图属性预测任务下，训练样本数据的构建方式和一般视觉、自然语言任务类似，输入的单个样本是一个完整的图，分布式训练可以在样本层面做拆分。具体方式如下例所示：
from torch.utils.data.distributed import DistributedSampler
from torch_geometric.loader import DataLoader
train_sampler = DistributedSampler(train_data)
train_loader = DataLoader(train_data, batch_size=16, sampler=train_sampler, num_workers=8, pin_memory=True)
而点属性预测和边链接预测则不同，训练样本一般是从一张大规模图中随机游走采样而成，特别是当需要考虑一些多步特征的情况，分布式训练所需的训练样本拆分就难以操作。
那么这两种任务就无法进行并行训练了呢？幻方AI就这个问题进行了调研和实验，总结了如下两种并行训练方式：
1. 分割子图
一种并行训练的方式是对大规模图进行图分割。
一般来说大规模图数据里，或多或少会存在一些社群分割，例如社交网络、学术网络等。在同一社群中节点彼此比较相似，链接性更高。因此，对于一个大规模图，可以将图根据连通性切割分为若干个子图，然后这些子图上进行分布式数据并行训练。
然而，这种方式需要进行数据的预处理，前期比较耗时。同时，图分割通过某一指标进行切分，实际上人为去除了某些社群间链接的特征，这在某些场景上来说，不可避免会对最终的落地效果造成影响。
2. 并行轮次
另外一种并行训练的方式，是并行轮次。
并行轮次的出发点是：既然不能在每 Epoch 上将训练样本拆分到多张卡上计算，那可以考虑让每张计算显卡单独计算一个 Epoch，合并梯度，减少训练轮次，进而提升训练效率。
以 DeepGCNs 的点预测任务为例：
for epoch in range(args.epochs):
train_parts = gdata.random_partition_graph(gdata.total_no_of_nodes, cluster_number=args.cluster_number)
train_data = gdata.generate_sub_graphs(train_parts, cluster_number=args.cluster_number)
...可以看到，在单卡上每 Epoch 都需要在这张大图 gdata 中随机划分出多个部分，然后依此构建多个子图，组成训练样本。这种方式在单卡训练上没有问题，然后扩展到多卡并行会启动多个进程，每轮每个进程都会单独随机采样出训练样本，这使得无法使用 DistributedSampler 在多卡对训练样本进行划分。
因此，索性不对训练样本进行划分，让每进程（或者说每张卡）单独完成一个 Epoch 的训练，然后再同步梯度。
值得注意的是，这种方式虽然简单易执行，但实际效果上是增加了 Batch Size，因此相应的需要增大一些学习率以保证训练效果。另一方面，每个进程都需要从保存在内存里的大图中采样出子图样本，因此对集群的性能要求相应比较高。
实验
这里我们通过实验来验证并行训练的效果。
数据集
幻方AI的数据集仓库集成了 OGB 大规模图公开数据集，通过 3FS 高速存储加速图数据的读取。具体的使用方式如下：
from hfai.datasets import OGB
gdata = OGB(data_name='ogbn-proteins')
gdata.get_data()
for epoch in range(args.epochs):
# training model ...如上例代码，我们加载 hfai 数据仓库中的 ogbn-proteins 数据，将这个大图加载入内存。
模型训练
我们测试“并行轮次”的方法，观察模型的收敛速度。首先我们测试单卡上，DeepGCNs 的收敛情况：
之后，我们将实验扩展到8卡并行，测试 DeepGCNs 的收敛：
从上面两次测试结果中可以看到，使用多卡训练下模型收敛的速度更快，因为在同一时间模型喂入了更多的子图数据，梯度同步获得的损失信息更多。这里需要注意，因为子图是通过随机采样出来的，因此在不同的进程中需要使用不同的随机算子，使得喂入的子图是不一样的。从实验可得，并行轮次可以有效提升训练效率，更快模型收敛，进而能直接节省训练的轮次，节省训练时间。
体验总结
综上，图神经网络的并行加速还是切实可行的。采用不同于常规分布式训练的方式，更多的计算资源可以被利用起以加速图神经网络的训练。
综合体验打分如下：
-
研究指数：★★★★
该模型着眼的是图数据场景，具有很广泛的研究应用场景。
-
开源指数：★★★★★
数据和代码都已开源，代码逻辑清晰可读性高。
-
门槛指数：★★★
数据规模中等，模型计算对算力的需求较高，需要多显卡并行加速训练。
-
通用指数：★★★★
模型是通用的图神经网络模型，在很多图数据上都可以适配应用。
-
适配指数：★★★★★
依赖简单，很容易与幻方AI的训练优化工具结合，提效明显。
幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力AI落地与价值创造，欢迎各方数据研究者与开发者们一同共建。