图是一种由边连接的节点组成的潜在的广泛网络，可以用来表达和查询数据之间的关系，如社会关系、金融交易、交通、能源网格和分子相互作用。随着研究人员收集更多的数据并构建这些图形图像，研究人员将需要更快、更有效的方法，以及更强的计算能力，以图神经网络(GNN)的方式对它们进行深度学习。

现在，麻省理工学院和IBM研究院的研究人员开发了一种名为SALIENT(采样、切片和数据移动)的新方法，通过解决计算中的三个关键瓶颈来提高训练和推理性能。这大大减少了大型数据集上gnn的运行时间，例如，包含1亿个节点和10亿个边。此外，该团队发现，当计算能力从1个图形处理单元(gpu)增加到16个时，该技术的扩展性很好。这项工作在第五届机器学习和系统会议上发表。

“我们开始关注当前系统在将最先进的机器学习技术用于图形到真正大的数据集时所面临的挑战。事实证明，还有很多工作要做，因为许多现有的系统主要在适合GPU内存的较小数据集上取得了良好的性能，”该研究的第一作者、麻省理工学院计算机科学和人工智能实验室(CSAIL)的博士后蒂姆·卡勒说。

专家们所说的庞大数据集，指的是整个比特币网络的规模，其中的某些模式和数据关系可能阐明趋势或欺诈行为。“区块链上有近10亿比特币交易，如果我们想要识别这样一个联合网络中的非法活动，那么我们就要面对这样一个规模的图表，”IBM研究院和MIT-IBM沃森人工智能实验室的高级研究科学家和经理、合著者陈杰(Jie Chen)说。“我们希望建立一个能够处理这种图形的系统，并允许处理尽可能高效，因为我们每天都希望跟上新数据生成的速度。”

Kaler和Chen的合著者包括Jump Trading的Nickolas Stathas孟'21，他开发了SALIENT作为其研究生工作的一部分;前MIT- ibm沃森人工智能实验室实习生、MIT研究生Anne Ouyang;麻省理工学院CSAIL博士后Alexandros-Stavros Iliopoulos;MIT CSAIL研究科学家Tao B. Schardl;以及麻省理工学院埃德温·西布里·韦伯斯特电气工程教授、MIT- ibm沃森人工智能实验室研究员查尔斯·e·雷瑟森。

对于这个问题，该团队采取了一种系统导向的方法来开发他们的方法:SALIENT, Kaler说。为了做到这一点，研究人员实现了他们认为重要的基本优化组件，这些组件适合现有的机器学习框架，如PyTorch几何和深度图库(DGL)，它们是用于构建机器学习模型的接口。斯塔萨斯说，这个过程就像更换引擎来制造更快的汽车。他们的方法设计适合现有的GNN架构，因此领域专家可以轻松地将这项工作应用到他们的特定领域，以加快模型训练，并在推断过程中更快地梳理出见解。该团队确定，诀窍是让所有硬件(CPU、数据链和GPU)始终处于忙碌状态:当CPU采样图并准备小批量数据，然后通过数据链传输时，更关键的GPU正在工作，以训练机器学习模型或进行推断。

研究人员首先分析了gnn常用的机器学习库(PyTorch Geometric)的性能，该库显示出可用GPU资源的利用率低得惊人。通过简单的优化，研究人员将GPU利用率从10%提高到30%，相对于公共基准测试代码，性能提高了1.4到2倍。这种快速基线代码可以在50.4秒内通过算法(一个历元)对大型训练数据集执行一次完整的遍历。

为了进一步提高性能，研究人员开始研究发生在数据管道开始的瓶颈:图采样和小批准备的算法。与其他神经网络不同，gnn执行邻域聚合操作，它使用图中其他邻近节点的信息计算关于一个节点的信息——例如，在社交网络图中，用户的朋友的朋友的信息。随着GNN层数的增加，网络中获取信息的节点数量可能会激增，超过计算机的极限。邻域抽样算法通过选择更小的随机节点子集来进行收集;然而，研究人员发现，目前的实现速度太慢，无法跟上现代gpu的处理速度。作为回应，他们确定了数据结构、算法优化等的组合，从而提高了采样速度，最终将采样操作提高了约三倍，将每个周期的运行时间从50.4秒提高到34.6秒。他们还发现，在推断过程中，以适当的频率进行采样，可以提高整体能源效率和性能，该团队指出，这一点在文献中被忽略了。

在以前的系统中，这个抽样步骤是一种多进程方法，这会在进程之间产生额外的数据和不必要的数据移动。研究人员通过创建一个具有轻量级线程的单一进程，将数据保存在共享内存中，从而使他们的SALIENT方法更加灵活。此外，显著利用了现代处理器的缓存，Stathas说，并行特征切片，从CPU核心缓存的共享内存中感兴趣的节点及其周围的邻居和边缘提取相关信息。这再次将整个周期运行时间从34.6秒减少到27.8秒。

研究人员解决的最后一个瓶颈是使用预取步骤在CPU和GPU之间流水线小批数据传输，这将在需要数据之前准备数据。该团队计算出，这将最大化数据链的带宽使用，并使该方法达到完美利用;然而，他们只看到了大约90%。他们在一个流行的PyTorch库中发现并修复了一个导致CPU和GPU之间不必要的往返通信的性能bug。修复了这个bug后，团队使用SALIENT实现了16.5秒的周期运行时间。

卡勒说:“我认为，我们的研究表明，细节决定成败。”“当你在训练一个图神经网络时密切关注影响性能的细节，你可以解决大量的性能问题。在我们的解决方案中，我们最终完全被GPU计算所困，而这正是这种系统的理想目标。”

在三个标准数据集ogbn-arxiv、ogbn-products和ogbn-papers100M上评估了SALIENT的速度，以及在多机设置下，使用不同级别的fanout (CPU将为GPU准备的数据量)，并跨越多个架构，包括最新的最先进的架构GraphSAGE-RI。在每一个设置中，SALIENT都优于PyTorch Geometric，最显著的是在大型ogbn-papers100M数据集上，包含1亿个节点和超过10亿个边。这里，它在一个GPU上运行的速度是最初为这项工作创建的优化基线的三倍;在配备16个gpu的情况下，SALIENT的速度增加了8倍。

虽然其他系统的硬件和实验设置略有不同，所以并不总是直接的比较，但SALIENT仍然优于它们。在达到类似精度的系统中，代表性的性能数字包括使用一个GPU和32个cpu用时99秒，使用1536个cpu用时13秒。相比之下，使用1个GPU和20个cpu时，SALIENT的运行时间为16.5秒，而使用16个GPU和320个cpu时仅为2秒。Kaler说:“如果你看看之前工作报告的底线数字，我们的16个GPU运行时间(2秒)比之前报告的其他数据快了一个数量级。”研究人员将他们的性能改进部分归因于他们在转移到分布式设置之前针对单个机器优化代码的方法。斯塔萨斯说，这里的教训是，对于你的钱来说，“在你开始扩展到多台计算机之前，更有意义的是有效地使用你的硬件，并将其发挥到极致”，这可以在成本和碳排放方面提供显著的节省，这可以与模型训练相结合。

这种新的能力现在将允许研究人员处理和深入挖掘越来越大的图表。例如，前面提到的比特币网络包含10万个节点;SALIENT系统能够处理1000倍(或3个数量级)大的图形。

Chen说:“未来，我们不仅要在现有的算法上运行这个图神经网络训练系统，我们已经实现了分类或预测每个节点的属性，我们还想做更深入的任务，比如识别图中的常见模式(子图模式)，(这)实际上可能对指示金融犯罪很有趣。”“我们还想在图中识别相似的节点，这些节点在某种意义上可能对应于金融犯罪中的同一坏人。这些任务将需要开发额外的算法，可能还需要神经网络架构。”

这项研究得到了麻省理工学院- ibm沃森人工智能实验室的支持，部分得到了美国空军研究实验室和美国空军人工智能加速器的支持。

注：本文由院校官方新闻直译，仅供参考，不代表指南者留学态度观点。