图神经网络系统的挑战

图神经网络借助图强大的数据表示能力和神经网络参数共享、局部连接和强大的拟合能力在社交网络分析、推荐系统、分子预测等多种领域取得了很好的效果,飞速增长的数据规模给图神经网络系统带来了可扩展性的问题。尤其是图中节点之间不规则的连接在基于采样的GNN训练中导致了节点的冗余访问,相较于图拓扑数据而言庞大的特征数据在不同的节点或设备之间的传输导致的通信效率低下,最终带来的表现就是GNN系统的高存算比。图不规则访问的特点也导致图划分、图采样加速存在困难,GNN系统的不同模块之间的组织和架构设计也对GNN系统的效率有很大影响。

本文将简要介绍现有GNN系统中的挑战和可能的解决方案。

减少数据冗余访问

在深度神经网络中,针对的数据均为欧几里得数据,也就是序列化的数据如文本、图像、音频等,这种场景下样本与样本之间是独立的,训练当前样本时无需其它样本的参与。与传统的深度神经网络不同,图神经网络数据集中样本和样本之间存在复杂的依赖性,在对当前样本进行训练时,必须同时提取依赖的其它样本,这就为数据访问带来了冗余性,在一个完整的epoch内,样本会被访问多次,这意味如果使用GPU训练的话,节点也可能会在CPU内存和设备内存之间重复传输多次。我们希望能够尽可能地减少特征数据的传输,因为它们的体积实在是太庞大啦。

探索模型并行

考虑到图神经网络的模型一般都比较小,所以我们谈及GNN模型的规模时,通常指的是数据集的规模。对于大规模的计算可以使用并行计算或分布式计算来加快整体计算速度,然而并行或分布式会带来一些额外的开销,并不能达到线性加速比,所谓线性加速比是指系统的计算效率随着节点数量的增加呈线性提升,一般情况下随着节点数量的增加,计算效率提升得越来越慢。

目前主流的图学习框架如DGL、PyG等,均使用数据并行的模式来进行单机多卡或多机多卡训练,即在每个GPU的设备内存中都维护一份完整的模型,将训练样本划分成多份分发给多个GPU进行训练,每个batch的数据训练完后,GPU之间进行梯度同步反向传播更新参数。数据并行的问题在于设备之间流动的是图数据,模型之间还需要进行通信来完成梯度同步,两者都会带来较大的延迟,影响训练效率。另外一种并行方法是模型并行,即将模型按照层或Tensor拆分分散在不同的设备上,P3提供了模型并行的新思路,它将数据分散固定在不同的设备上,分割的模型在设备间流动,由于GNN中模型的体积远小于图数据集本身,这大大优化了分布式训练的通信瓶颈。

更高效的算法

Semi-Supervised Classification with Graph Convolutional Networks中提出GCN时,训练还采用full batch的模式,这也决定了GCN是直推式的,无法将训练好的模型应用到新的节点上,同时full batch的训练所需要的内存与图规模呈正比,在有限的内存空间下无法训练大图。

Inductive Representation Learning on Large Graphs使用了基于采样的训练方式,从训练节点的K阶邻居节点中随机选择若干个节点,并通过聚合函数和更新函数来得到目标节点的新的表示。注意GraphSAGE与GCN不同,后者基于谱域理论,训练时需要用到全图数据,而前者基于空域理论,模型本质上在学习一组聚合邻居节点特征的函数,这与全图结构是无关的,因此GraphSAGE可以泛化至新的节点上。

GraphSAGE使用node-wise的采样思想,在每层中采样特定数量的邻居,随着层数的增长,采样节点数指数增长,这就是邻域爆炸的问题。cluster-gnn在聚类的基础上进行采样,属于subgraph-wise的采样思想,在一个子图中随机选取特定数量的节点来聚合更新目标节点。FastGNN等则使用了layer-wise的思想,在每层中采样固定数量的节点,这样采样子图的大小随层数增加呈线性增长而不是指数增长,一定程度上缓解了邻域爆炸的问题。然而采样算法对方差和偏差是有要求的,我们希望算法的方差和偏差尽可能地小,这样才能保证模型的准确率和收敛性。探索能适应于大图训练的快速高效的算法是GNN领域中的一大挑战。

模块化

一个完备的图神经网络系统应当包含了若干个功能模块,模块自身高内聚、模块之间低耦合,同时要具备良好的可扩展性。数据模块、数据加载器模块、采样模块、特征提取引擎模块、模型模块。

存外计算

存外计算是在单机上训练超大规模图的另一种方式,将当前训练所需要的数据加载到内存中,而将大部分不需要的数据储存在外存中。这种做法需要实现一个统一内存外存的接口,并在内存外存间设计数据交换策略,来保证大部分数据在内存中命中,掩盖磁盘的IO延迟,代表性的工作有Marius。此外也可以利用GPU Direct Storage、Nvm等技术和设备来进行扩展和加速。

分区算法

图分区算法主要应用在分布式训练环境中,原因一是图的规模太大,单个节点已经无法存储全部的数据,必须划分后分别储存在多个节点上,原因二是分布式环境中节点之间的通信开销是一个很大的瓶颈,图的划分要尽可能得减少计算的跨分区访问。

分区算法应该有高效的划分效率。图的邻接矩阵需要$O(n^2)$的空间复杂度,图的各种算法也具有较高的时间复杂度,如Dijkstra需要$O(n^2)$的时间复杂度,Floyed需要$O(n^3)$的时间复杂度,使用BFS遍历一张图也需要$O(|V|+|E|)$的时间复杂度。因此分区算法应该具有较低的时间复杂度,或者可以通过将并行计算来加速。

为了保证多个节点之间的梯度同步,分区之间还要做到负载均衡,即每个分区内拥有的训练节点要保持平衡。

GPU采样

采样属于计算密集型任务,在典型的GNN系统中,均在CPU上完成采样,而CPU还要负责特征提取和任务调度等其他工作,尤其在多卡场景下,多个GPU对CPU存在竞争。图采样看起来是一个很适合并行计算的场景,但由于图访问的不规则性,在GPU上进行图采样很难利用GPU的资源,Nextdoor旨在解决这个问题,现在DGL这样的通用图神经网络框架也支持使用GPU采样了。GPU采样大大提高了采样的速度,但为整个训练流程引入了更加复杂的因素,要知道特征提取始终由CPU来完成,batch传送到GPU采样后,采样子图再传送到CPU内存中,完成特征提取后再次传输到GPU内存中进行训练。我们需要仔细地考虑GPU采样可能带来的延迟,这对于一些实时性要求较高的场景可能非常关键。

总结

以上提到的几个方面只是GNN系统中最常见的优化角度,实际上在消息传递框架、对异构图和动态图的支持、缓存策略等方面也面临着诸多问题。

暂无评论

发送评论 编辑评论


				
上一篇
下一篇