图神经网络实战——大规模的学习与推理

本章内容包括：

• 处理小型系统中数据过载的策略
• 识别需要扩展资源的图神经网络问题
• 缓解大数据带来的问题的七种稳健技术
• 使用PyTorch Geometric扩展图神经网络并应对可扩展性挑战

在我们对图神经网络（GNNs）的探索过程中，我们介绍了关键的架构和方法，但示例大多限于相对较小规模的问题。我们这样做的原因是希望让你能够轻松访问示例代码和数据。

然而，深度学习中的现实问题往往并不像这样简单地被打包成一个个小问题。现实场景中的一个主要挑战是，当数据集足够大，无法完全装入内存或压倒处理器时，如何训练GNN模型。

在探讨可扩展性挑战时，清晰地理解GNN训练过程至关重要。图7.1重新展示了我们熟悉的训练过程可视化图。其核心是，GNN的训练围绕着从数据源获取数据、处理这些数据以提取相关的节点和边特征，然后使用这些特征训练模型。随着数据规模的增长，每一个步骤可能都会变得更加资源密集，因此，我们将在本章中探索的可扩展策略变得尤为重要。

在深度学习开发项目中，考虑到训练和部署中的大规模数据可以决定项目的成败。对于在紧迫期限内工作并且有苛刻需求的机器学习工程师来说，他们没有奢侈的时间来花费数周进行长时间的训练流程或修复由处理器超载引发的错误。通过提前规划并预防规模问题，可以避免这些时间上的浪费。

在本章中，你将学习如何处理当数据过大而无法适应小型系统时所出现的问题。为了描述规模问题，我们关注三个指标：处理或训练过程中的内存使用情况、训练一个周期所需的时间以及问题收敛所需的时间。我们将解释这些指标，并指出如何在Python或PyTorch Geometric（PyG）环境中计算它们。

本章的重点是从小规模开始进行扩展，从单一机器进行优化。尽管本书的主要关注点并不在数据工程或大规模解决方案的架构上，但一些讨论的概念在这些上下文中可能是相关的。为了解决规模问题，本章解释了七种方法，可以一起使用或单独使用：

1. 选择和配置处理器（第7.4节）
2. 使用稀疏与密集的表示方式来处理数据集（第7.5节）
3. 选择GNN算法（第7.6节）
4. 基于数据采样进行批量训练（第7.7节）
5. 使用并行或分布式计算（第7.8节）
6. 使用远程后端（第7.9节）
7. 粗化图（第7.10节）

为了说明如何在实际中做出这些方法的决策，本章提供了例子或迷你案例。我们将通过虚构公司GeoGrid Inc.（以下简称GeoGrid）的案例，跟踪公司如何处理与大数据相关的实际问题。

此外，在第3章中遇到的亚马逊产品数据集，也用于展示这些方法。在该数据集中，使用图卷积网络（GCN）和GraphSAGE进行节点分类。本章相关方法的示例代码可以在本书的GitHub仓库中找到。

本章与之前的章节有所不同。之前的章节主要集中在通过一个或两个示例来说明一系列概念，而规模问题的独特性质意味着我们需要探索各种方法，并为每种方法提供简短的示例。因此，本章的各节可以在第7.3节之后按任意顺序阅读。

我们将首先回顾第3章中的亚马逊产品数据集，并介绍GeoGrid。接着，我们将讨论如何描述和衡量规模，重点关注这三个指标。最后，我们将更详细地介绍每种方法，并在适当的地方提供代码。

注：本章的代码可以在GitHub仓库中的Notebook形式找到（mng.bz/QDER）。Colab链接和本章的数据也可以在相同的位置访问。

7.1 本章中的示例

在本章中，使用了两个案例来说明各种概念。我们将使用第3章中的亚马逊产品数据集。我们将用这个数据集来展示代码示例，代码可以在GitHub仓库中找到。其次，我们将使用一个名为GeoGrid的虚构公司的小案例，来阐明使用本章介绍的方法的指导原则和实践。

7.1.1 亚马逊产品数据集

本小节将重新介绍第3章中的数据集及其训练过程。首先，我们回顾一下数据集，然后介绍用于训练的硬件配置。最后，作为接下来各节的前奏，我们重点介绍在第3章中应用的一些方法，以适应数据集的大小。这个数据集将在接下来的GitHub代码示例中被广泛使用。

在第3章中，我们研究了使用两种卷积GNN进行的节点分类问题：GCN和GraphSAGE。为此，我们使用了包含共同购买信息的亚马逊产品数据集，这个数据集广泛用于说明和基准测试节点分类[2]。该数据集（也称为ogbn-products）由一组产品节点组成，这些节点通过在同一交易中被购买而相连，如图7.2所示。每个产品节点都有一组特征，包括其产品类别。ogbn-products数据集包含250万个节点和6190万个边。关于该数据集的更多信息，概述在表7.1中。

注：有关此数据集及其来源的更多细节，以及GCN和GraphSAGE的内容，请参阅第3章。

表7.1 ogbn-products数据集的概要特征

特征	数值
节点数量	250万
边数量	6190万
平均节点度	51
类标签数量	47
节点特征维度	100
压缩数据大小（GB）	1.38

在第3章实现的代码中，我们使用了配置如下的Colab实例：

• 存储：56 GB HDD
• 两个CPU：2核Xeon 2.2GHz
• CPU内存：13 GB
• 一个GPU：Tesla T4
• GPU内存：16 GB

虽然我们稍后会详细讨论，但我们已经识别出三个可能影响我们是否会因数据过大而遇到问题的因素。一个显然是数据集本身的大小——不仅仅是存储中的原始解压大小，还包括其表示方式，这会影响在进行处理和训练时的工作大小（第7.5节将详细讨论）。第二个因素是硬件的存储和内存容量（第7.4节）。最后，GNN训练算法的选择——例如GraphSAGE——将显著影响计算需求，特别是在时间和内存限制方面（第7.6节）。

在实现第3章中的示例时，我们确实遇到了问题，根本原因是数据集的大小。我们在该章的重点是展示算法，因此并没有指出这一点，而是悄悄地使用了其中的一种方法来缓解这个问题。具体来说，我们使用了数据集的最优表示方式（稀疏表示而非密集表示）。

7.1.2 GeoGrid

在本章的学习过程中，我们将通过一个虚构但具有代表性的技术公司——GeoGrid——的例子，来探讨该领域中的挑战和机遇。GeoGrid是一家地理空间数据分析和建模公司。通过使用图神经网络（GNN）等先进技术，公司提供从交通预测到气候变化规划等多种问题的解决方案。作为一家在竞争激烈领域中的初创公司，GeoGrid经常面临可能决定公司成败的关键技术决策，尤其是在争夺大型政府项目时。

GeoGrid将被用于探讨与规模问题相关的一系列概念和技术决策。无论是团队在讨论不同机器学习架构的优缺点，考虑在多个GPU上使用分布式数据并行（DDP）训练，还是制定如何为大规模数据集扩展算法的策略，GeoGrid的故事都为本章讨论的理论和方法提供了现实世界的背景。

在下一节中，我们将提供一个框架来判断和描述规模问题。接着，我们将总结解决这些问题的方法。最后，我们将详细探讨这些方法。

7.2 规模问题的框架

在我们深入探讨解决方案之前，首先定义一下扩展所带来的挑战。本节概述了数据规模问题的根本原因及其症状。然后，强调了在识别、诊断和解决这些问题时至关重要的指标[1, 3]。

从机器资源的角度来看，开发过程可分为三个阶段。以下三阶段中，本章的重点将放在预处理和训练阶段：

• 预处理——将原始数据集转换为适合训练的格式
• 训练——通过将训练算法应用于预处理后的数据集来创建GNN模型
• 推理——从训练好的模型中生成预测或其他输出

7.2.1 根本原因

简单来说，规模问题出现在训练数据过大，超出我们系统的承载能力时。确定数据何时变得问题化是复杂的，这取决于多个因素，包括硬件能力、图的大小以及时间和空间的限制。

硬件速度和容量

合适的系统必须能够通过其内存容量和处理速度支持预处理和训练过程。内存不仅要支持图的大小，还需要容纳用于实现转换和训练算法的数据。处理速度应该足够快，以便在合理的时间内完成训练。

我们在编写本书时假设你可以访问如Google Colab和Kaggle等免费的云资源，或者至少有一台带有GPU处理器的本地资源。当这些资源不足时，升级硬件配置可能是一个选择（如果资源允许的话）。对于最大规模的企业图，使用计算集群是不可避免的。我们将在第7.4节中更详细地探讨计算硬件。

图的大小

从根本上说，我们可以通过节点和边的数量来大致了解规模及其对训练解决方案的影响。理解这些特征可以帮助我们估算一个算法处理图所需的时间。此外，保存结构信息的数据表示将影响数据的大小。

除了结构信息之外，节点和边还可以包含涵盖一个或多个维度的特征。通常，节点和边特征的大小可能比图的结构信息更大。

对于GNN来说，小型、中型和大型图的确切定义是有一定上下文的。这取决于特定的问题领域、硬件和可用的计算资源。在本书写作时，以下是一个大致的分类：

• 小型图——这些图可能包括数百到几千个节点和边。通常可以在标准硬件上处理，而无需专门的资源。
• 中型图——这一类别可能包括数万个节点和边。中型图的复杂性可能需要更复杂的算法或硬件（例如GPU）才能高效处理。
• 大型图——大型图可以包括数十万个到数百万（甚至数十亿）个节点和边。处理这样的图通常需要分布式计算和为可扩展性设计的专门算法。

算法的时间和空间复杂度——时间和空间复杂度指的是运行算法所需的计算和内存资源。这直接影响处理速度、内存使用和效率。理解这些复杂度有助于做出有关算法选择和资源分配的明智决策。高时间复杂度可能导致运行时间较长，影响模型训练进度。高空间复杂度则可能限制GNN能够处理的数据集大小，从而影响处理大型复杂图的能力。我们将在第7.6节中进一步探讨这一点。

7.2.2 症状

规模问题的根本原因表现为几种方式。一种常见的问题是处理时间过长，这通常发生在较大的数据集需要更多的计算能力和时间来处理时。较慢的算法可能会增加训练模型所需的时间，从而使得模型的迭代和改进变得困难。然而，何时视为时间过长，取决于具体问题的要求。对于需要每周提供结果的问题，几小时可能是可以接受的，但如果模型需要全天候重新训练，那么几小时可能就显得太长。同样，如果处理时间很长，计算成本也可能会迅速增加，特别是当需要使用大型机器来运行模型时。

另一个问题是内存使用达到或超过容量，这通常发生在大型数据集消耗了大量内存时。如果数据集太大，无法适应系统的内存，就可能导致系统变慢甚至崩溃。

最后，当算法和系统设置无法应对数据集增大时，就会出现无法扩展到更大数据集的问题。确保在时间和空间上的效率对于保持系统的有效性和可扩展性至关重要。

7.2.3 关键指标

为了理解可扩展性问题，运行关键性能指标的实证分析是非常有帮助的。这些指标包括内存使用、每个周期的时间、FLOPs和收敛速度，具体描述如下：

内存使用——内存使用量（单位：GB），尤其是可用的RAM或处理器内存，在决定你能够训练的模型的大小和复杂性方面起着重要作用[4, 5]。这是因为GNN需要将节点特征、边特征和邻接矩阵存储在内存中。如果图非常大，或者节点和边的特征维度很高，模型将需要更多的内存。

PyTorch和Python中有多个模块可以进行内存分析。PyTorch有一个内置的分析器，可以单独使用或与PyTorch Profiler Tensorboard插件结合使用[4]。还有一个torch_geometric.profile模块。此外，Colab和Kaggle上托管的云笔记本可以实时可视化每个处理器的内存使用情况。

在我们仓库中的代码示例中，我们使用了两个库来监控系统资源：psutil（Python系统和进程工具库）和pynvml（NVIDIA管理库的Python绑定）。psutil是一个跨平台的工具，提供检索系统利用率（CPU、内存、磁盘、网络、传感器）、运行中的进程和系统运行时间的信息接口。它特别适用于实时系统监控、分析和限制进程资源。以下是psutil在代码中的使用示例：

import psutil

def get_cpu_memory_usage(): 
    process = psutil.Process(os.getpid()) 
    return process.memory_info().rss

在这个示例中，psutil.Process(os.getpid())用于获取当前进程，memory_info().rss检索驻留集大小，即进程占用的RAM内存部分。

除了psutil，pynvml是一个用于与NVIDIA GPU交互的Python库。它提供GPU状态的详细信息，包括使用情况、温度和内存。pynvml允许用户通过编程方式检索GPU统计信息，是管理和监控机器学习及其他GPU加速应用中GPU资源的重要工具。以下是pynvml在代码中的使用示例：

import pynvml

pynvml.nvmlInit()

def get_gpu_memory_usage(): 
    handle = pynvml.nvmlDeviceGetHandleByIndex(0) 
    info = pynvml.nvmlDeviceGetMemoryInfo(handle) 
    return info.used

在此，pynvml.nvmlInit()初始化NVIDIA管理库，pynvml.nvmlDeviceGetHandleByIndex(0)检索索引为0的GPU句柄，pynvml.nvmlDeviceGetMemoryInfo(handle)提供GPU内存使用的详细信息。

psutil和pynvml都在我们的示例中用于提供有关预处理和训练过程的性能特征，详细显示系统和GPU资源的利用情况。

每个周期的时间——每个周期的时间（又称“每周期秒数”，因为该指标的单位通常是秒）指的是完成对整个训练数据集的遍历所需的时间。这个因素受到GNN的大小和复杂性、图的大小、批次大小以及可用计算资源的影响。每个周期时间较短的模型更受欢迎，因为它可以进行更多的迭代和更快的实验。PyTorch或PyG提供的分析器也可以用于这种测量。

在提供的代码中，每个周期所用时间是通过计算周期开始和结束时间的差值来衡量的。在每个周期开始时，使用start_time = time.time()记录当前时间。然后训练模型1个周期，完成后再次使用end_time = time.time()记录当前时间。周期时间（即完成1个训练周期所花费的时间）然后计算为结束时间和开始时间的差值（epoch_time = end_time - start_time）。这给出了训练1个周期所花费的确切时间，包括训练过程中的所有步骤，如前向传播、损失计算、反向传播和模型参数更新。

FLOPs——浮点运算（FLOP，即浮点运算数，不要与浮点运算每秒数（FLOP/s）混淆[6, 7]）计算训练一个模型所需的浮点运算数量。这可以包括矩阵乘法、加法和激活等操作。对于我们的目的，FLOPs的总数提供了训练GNN的计算成本估算。

FLOPs在执行时间上并不完全相同。这个差异源于几个因素。首先，涉及的操作类型可以极大地影响计算成本：简单的加法和减法通常较快，而更复杂的操作，如除法或平方根计算，通常较慢。其次，FLOPs的执行时间可能会根据所使用的硬件而有所不同。某些处理器针对特定类型的操作进行了优化，专用硬件如GPU可能比CPU更有效地处理某些操作。此外，算法的结构也会影响FLOPs的执行效率；可以并行化的操作可能在多核系统上处理得更快，而依赖于先前结果的顺序操作可能总体上花费更多时间。尽管执行时间有所不同，但给定算法所需的FLOPs总数保持不变。

在编写时，虽然有些外部模块可以分析PyTorch操作，但它们与PyG模型和层不兼容。文献中的一些努力依赖于自定义编程。

在我们GitHub上的代码示例中，我们通常使用thop库来估算每个周期训练神经网络时的FLOPs。以下是FLOPs计算的简短示例：

from thop import profile  #1

input = torch.randn(1, 3, 224, 224) 
macs, params = profile(model, inputs=(input, )) 
print(f"FLOPs: {macs}")

在此，profile函数被调用，模型和一个示例输入批次作为参数传递。它返回一个前向传播所需的总FLOPs和参数。在这个上下文中，FLOPs衡量的是总操作数，而不是每秒操作数。

FLOP是一个有用的指标，用于大致了解模型的计算需求和复杂性，特别是与其他指标结合使用时，能够全面理解性能。

收敛速度——收敛速度（单位：秒或分钟）指的是模型在训练过程中学习或达到最优状态的速度。收敛速度受到模型复杂性、学习率、使用的优化器和训练数据质量等因素的影响。较快的收敛通常是可取的，因为它意味着模型需要较少的周期来达到最优状态，从而节省时间和计算资源。

与内存和每个周期时间分析一样，PyTorch和PyG的分析器也可以用于测量收敛所需的时间。

在我们的代码示例中，收敛时间是通过测量完成指定周期数的训练所需的时间间隔来计算的。在训练过程开始时，使用convergence_start_time = time.time()记录开始时间。然后，模型经历多个周期的训练，每个周期涉及前向传播、损失计算、反向传播和参数更新等步骤。完成所有周期后，再次记录当前时间，并通过从最终时间戳中减去convergence_start_time来计算收敛时间。这个convergence_time给出了模型在所有周期训练结束时所花费的总时间，提供了模型在时间上的效率和性能的见解。收敛时间越短，模型学习并达到令人满意的性能水平的速度越快，前提是保持良好的学习质量。

这四个因素之间的平衡取决于具体项目的约束条件，如可用的计算资源、项目时间表以及数据集的复杂性和大小。对于这些指标的现实基准测试，Chiang [8]很好地利用这些指标对他提出的GNN、ClusterGCN和基准GNN进行比较分析。在了解了规模问题的构成，以及如何基准和测量这些问题后，我们将转向可以缓解这些挑战的方法。

7.3 解决规模问题的技术

正如我们在前一节中概述的，当数据量变得庞大时，我们必须应对与内存限制、处理时间和效率相关的问题。为了应对这些挑战，拥有一套策略工具箱变得至关重要。在接下来的几节中，我们将介绍一系列旨在提供灵活性和控制训练过程的方法。这些策略从硬件配置到算法优化不等，旨在适应不同的场景和需求。这些方法来自深度学习和图神经网络在学术界和工业界的最佳实践。

7.3.1 七种技术

首先，我们从三个基本选择开始，这些选择可以提前规划，并在项目过程中进行重新配置。为了准备你的项目，请选择以下内容：

• 硬件配置——这些选择涵盖了处理器类型、处理器的内存配置，以及是使用单台机器/处理器还是多台机器。
• 数据集表示——PyG支持密集和稀疏张量。在处理大图时，将数据从密集表示转换为稀疏表示可以显著减少内存占用。你可以使用PyG的torch_geometric.utils.to_sparse函数将密集邻接矩阵或节点特征矩阵转换为稀疏表示。
• GNN架构——某些GNN架构设计上就具备了计算高效且可扩展的特点，适用于大规模图。选择一种能够良好扩展的算法，可以显著减轻规模问题。

鉴于这三个类别的选择，如果问题超出了系统的处理能力，以下是可以用来缓解问题的技术：

• 采样——你可以在每次训练迭代中，从整个大图中采样出一部分节点或子图，而不是在整个大图上训练。这种做法的复杂性（加入采样和批处理例程）可以通过内存效率的提升来弥补。为了进行节点或图的采样，PyG提供了torch_geometric.sampler和torch_geometric.loader模块的功能。
• 并行性和分布式计算——你可以使用多个处理器或机器集群，通过在训练过程中将数据集从一台机器扩展到多台机器，减少训练时间。根据实施方式，可能需要一些开发和配置的额外开销。
• 使用远程后端——你可以将训练图数据集完全存储在后端数据库中，而不是存储在内存中，并在需要时加载小批量数据。这种方法的最简单例子是将数据存储在本地硬盘上，并迭代地从中读取小批量数据。在PyG中，这种方法称为远程后端。这是PyG中的一种相对较新的方法，已有一些示例但不多。截止写作时，两家数据库公司已经开发了一些对PyG远程后端功能的支持。此方法需要最多的开发和维护工作，但它在缓解大数据问题方面最为有效。
• 图粗化——图粗化技术用于在尽可能保留其基本结构的同时减少图的大小。这些技术将节点和边聚合，创建出一个比原始图更粗的版本。PyG为此提供了图聚类和池化操作。缺点是，你必须小心确保粗化后的图能够真实地代表原始图，并且对于监督学习，你需要决定如何合并目标。

训练GNN时的规模问题是一个多方面的问题，需要深思熟虑的方法。通过应用多种杠杆，如硬件选择、优化技术、内存管理和架构决策，你可以根据具体的需求和约束调整过程。

7.3.2 一般步骤

在本节中，我们提供了一些规划和评估项目时考虑规模问题的一般指导。以下是一般步骤：

• 规划阶段

  • 预见硬件需求——提前熟悉可用的硬件选项。许多在线和本地系统都有已发布的配置。
  • 了解你的数据——清楚每个机器学习生命周期阶段的数据集大小。
  • 内存与数据的比例——作为经验法则，内存容量应理想地是数据集大小的4到10倍。

• 基准测试阶段

  • 建立基准——使用具有代表性的数据集对这些指标进行基准测试。这些初步数据可以作为基础，用来预测训练和实验的时间表。
  • 训练的指标——监控并衡量关键指标，如内存利用率、每周期时间、每秒浮点运算次数（FLOP/s）和收敛时间。
  • 故障排除——如果遇到挑战且缺乏硬件升级的资源，可以考虑实现本章中详细介绍的策略，以绕过硬件限制。

现在我们已经了解了规模问题、衡量这些问题的指标以及一套缓解这些问题的技术，接下来让我们更详细地探讨这些方法。

7.4 硬件配置的选择

本节探讨选择和调整硬件配置来解决规模问题。首先，我们将回顾硬件配置的一般选择，然后广泛概述相关系统和处理器的选择。给出这些选项的指导方针和建议。本节最后将介绍第一个GeoGrid的小案例研究。

7.4.1 硬件选择类型

有多种硬件配置可用于训练GNN，每种配置都旨在满足不同的需求并优化性能：

• 处理器类型——PyTorch提供了在不同类型的处理器上运行的灵活性，包括中央处理单元（CPU）、图形处理单元（GPU）、神经处理单元（NPU）、张量处理单元（TPU）和智能处理单元（IPU）。虽然CPU是普遍存在的，可以处理大多数常规任务，但GPU具备并行处理能力，专为处理密集计算任务设计，非常适合训练大规模神经网络模型。TPU是针对机器学习任务的定制加速器，能够提供更强大的计算能力，但它们的可用性可能会受到限制。有关更多详细信息，请参见下一小节。NPU（专为手机、笔记本和边缘设备运行神经网络工作负载而设计的处理器）和IPU（为需要大规模数据处理的高度并行工作负载设计）是重要的处理器类别。目前，PyTorch仅支持Graphcore IPU。
• 内存大小——每种处理器类型都配备了相应的RAM，RAM的大小在决定系统能够处理的工作负载规模方面起着关键作用。充足的RAM确保顺利的模型训练，尤其对于那些需要处理大量数据或具有复杂架构的网络。
• 单一GPU/TPU与多个GPU/TPU——对于那些幸运地可以访问多个GPU或TPU的人来说，多个GPU或TPU能够显著加快训练速度。PyTorch提供了DistributedDataParallel模块，利用多个GPU或TPU并行训练模型。这意味着可以将计算负载分配到多个设备上，从而加速迭代和模型收敛。
• 单机与计算集群——有时训练需求不仅仅需要单台机器，可能需要整个集群。在此上下文中，集群指的是由多台机器组成的集合，每台机器配备了不同的计算、内存和存储资源。如果你有这样的资源，PyTorch的DistributedDataParallel模块依然是小规模集群的首选工具。在这种情况下，它允许你将训练过程扩展到整个集群，在处理特别大的模型或海量数据集时尤为宝贵。

随着硬件能力的扩展——从单个处理器到多个设备，再到整个集群——规划、设置和管理的复杂性也在增加。根据任务需求和可用资源做出明智的决策，可以使这条路更加顺畅和高效。正如引言中所强调的，本章将重点讨论单机优化。

7.4.2 处理器和内存大小的选择

在转向硬件考虑的话题时，理解用于训练GNN的主要硬件选项非常重要：CPU、GPU、NPU、IPU和TPU。本节将简要概述每种类型的硬件，并提供它们应用的指导方针。这些要点概述在表7.2中。

中央处理单元（CPU） ——CPU在通用计算任务中表现优异，从数据预处理到模型训练都能胜任。然而，它们并未针对专门的深度学习任务进行优化，这可能影响其速度和效率。优点是，与其他硬件选项相比，CPU通常更具性价比，适用于更广泛的用户群体。

图形处理单元（GPU） ——GPU专为需要并行计算能力的任务而设计。从本书到目前为止的内容，你知道它们常常作为训练GNN的首选硬件，特别是在使用旨在最大化GPU并行性的库（例如PyG）时。本书中的大多数示例都在Colab平台上的NVIDIA GPU上运行，包括Tesla T4、A100和V100。

张量处理单元（TPU） ——TPU是一种专用选择，由Google构建以加速机器学习计算。它们提供快速的计算速度，并且可以具有成本效益。然而，它们的适用范围可能有限，因为它们是专有技术，主要与Google Cloud和TensorFlow兼容，并且可能不完全支持PyTorch。

神经处理单元（NPU） ——AMD和Intel都拥有NPU产品线，并配有可以与PyTorch集成的加速库。NPU是专为并行处理任务设计的硬件，类似于TPU。虽然GPU最初是为图形处理设计的，但它们通常包含专门用于机器学习任务的电路。NPU将这些电路作为专用单元，提高了效率和性能。Apple通常在其大多数笔记本和电脑中提供类似的专用单元（称为Apple Neural Engine [ANE]）。

智能处理单元（IPU） ——这些是专门的电路芯片，旨在并优化深度学习任务。IPU由Graphcore开发，专注于图计算。它们非常适合基于图的GNN模型，因为它们允许在消息传递过程中按需并行化独立任务。IPU与PyTorch和PyG兼容，但需要重新编写某些任务。其他设计非常大且强大的专用芯片的公司包括Cerebras和Groq。

配置考虑因素——选择硬件时，必须考虑内存限制，因为GNN通常因图数据的独特结构而数据密集。硬件选择还会影响训练和推理的速度。因此，权衡成本与性能之间的取舍是至关重要的，这应根据项目的具体需求量身定制。

选择用于PyTorch训练GNN的硬件时，主要需要考虑的因素包括处理器类型（例如CPU、GPU或TPU）、可用内存以及预算限制。表7.2中整理了这些考虑因素，便于快速参考。

表7.2 处理器选择的优缺点

硬件	推荐工作负载	优点	缺点
CPU	预处理	适合数据收集和预处理	相较于GPU和TPU训练速度较慢
GPU	训练	由于并行处理，适合训练	比CPU昂贵
TPU	预处理和训练	对深度学习任务计算速度更快且具有成本效益	需要特定的软件基础设施
NPU	训练	专为深度学习优化，尤其适用于设备端AI应用，减少对云服务的依赖	仅限于特定的AI工作负载，主要是神经网络任务
IPU	训练	尤其适用于基于图的任务，如GNNs	与NPU相比可能更复杂，需要更多编程和优化

需要注意的是，某些处理器类型在机器学习生命周期中的特定步骤中表现突出：

• 数据收集和预处理——CPU通常足以处理这些步骤。通常它们可以高效地处理各种任务，而不需要专门的硬件。然而，根据我们的经验，对于一些内存密集型、长时间的预处理步骤，如果有TPU可用，TPU的性能会更好。
• 模型训练——通常这是生命周期中计算最密集的部分，GPU通常是最佳选择。GPU专为并行处理而设计，加速神经网络的训练。特别是对于GNNs，它们经常涉及图中的多个节点和边的计算，GPU能够提供良好的加速效果。当TPU可用时，可能会提供更高的性能优势。
• 模型评估和推理——对于评估和推理，选择CPU还是GPU取决于具体的应用场景。如果成本效益更重要，CPU可能是更好的选择。TPU具有高计算速度和成本效益，对于大规模部署可能是一个不错的选择，但它们的使用范围相比于CPU和GPU更为有限。

注意，最佳的处理器选择可能会根据项目的具体需求而有所不同，例如模型的复杂性、数据集的大小、使用的平台以及可用的预算。本节以虚构公司GeoGrid的一个示例结束。

示例

Smith博士在GeoGrid工作，GeoGrid是一家领先的地图公司，参与一个使用GNN分析不同城市之间传染病传播的研究项目。她的数据集包含来自10,000个连接的城镇（节点）数据，每个城镇大约有1,000个节点特征。该数据集的大小为10GB。以下概述了使用GNN分析该项目时需要准备的不同步骤：

规划阶段

• 预见硬件需求——Smith博士回顾了她所在大学的计算资源，发现他们可以访问GPU和CPU，但TPU的数量目前有限。
• 了解数据——Smith博士估算她的数据集总大小约为10GB。通过探索性数据分析，她确定她的数据是稀疏的。
• 内存与数据的比例——她遵循经验法则，确保内存容量应为数据大小的4到10倍，因此她推断理想情况下需要至少40GB到100GB的内存。

基准测试阶段

• 使用数据集的子集，Smith博士基准测试了GPU和CPU上的数据预处理时间和模型训练时间。她发现使用GPU进行模型训练时，速度显著加快，正如预期的那样，但CPU在数据预处理方面表现相对较好。她决定将CPU用于预处理，GPU用于模型训练。

故障排除

• 在调查频繁的系统崩溃和内存错误原因时，Smith博士意识到她当前的GPU没有足够的内存来处理更大的图数据。由于当时没有更大内存的设备可以使用，她决定使用子图采样方法（详见第7.7节），使数据更适合当前硬件。

通过这个示例，我们看到了解数据集和可用资源、进行基准测试以设定期望值以及故障排除以在限制条件内找到解决方案的重要性。接下来，我们将研究如何选择数据表示方式。

7.5 数据表示的选择

根据输入图的特性，在PyG中如何存储和表示它们将影响时间和空间的约束。在PyG中，主要的数据类torch_geometric.data.Data和torch_geometric.data.HeteroData可以用两种格式来表示图：稠密格式和稀疏格式。在PyG中，稠密和稀疏表示的区别在于图的邻接矩阵和节点特征在内存中的存储方式。稠密表示具有以下特点：

• 整个邻接矩阵存储在内存中，包含零和非零元素，使用一个大小为N × N的二维张量，其中N是节点的数量。
• 节点特征存储在一个大小为N × F的密集二维张量中，其中F是每个节点的特征数。 这种表示方式内存占用较高，但当图是稠密的（即大多数节点之间有连接时），可以提供更快的计算速度；也就是说，邻接矩阵中非零元素的百分比较高，如附录A所述。

另一方面，稀疏表示具有以下特点：

• 邻接矩阵以稀疏格式存储，例如COO（坐标）格式，它只存储非零元素的索引和值。
• 节点特征可以存储为稀疏二维张量或字典，字典将节点索引映射到它们的特征向量。 这种表示方式在内存使用上更为高效，特别是当图是稀疏的（即图中只有少数节点相互连接时）；也就是说，邻接矩阵中非零元素的百分比较低，如附录A所述。然而，对于某些任务而言，稀疏表示可能会导致比稠密表示更慢的计算。

注：要了解稀疏或稠密格式之间的区别以及图是稀疏还是稠密的特性，请参见附录A第A.2节。

在PyG中，有两种方法可以将稠密数据集转换为稀疏表示：使用内置函数或手动执行转换：

• torch_geometric.transforms.ToSparseTensor——PyG中的这个转换函数可以将稠密的邻接矩阵或边索引转换为稀疏张量表示。它使用COO（坐标）格式构造一个稀疏的邻接矩阵。你可以应用这个转换到你的数据集，将稠密表示转换为稀疏表示：

  from torch_geometric.transforms import ToSparseTensor
  
  dataset = YourDataset(transform=ToSparseTensor())
  

• 手动转换——你也可以使用PyTorch或SciPy的稀疏张量功能手动将稠密的邻接矩阵或边索引转换为稀疏表示。你可以创建一个torch_sparse.SparseTensor或scipy.sparse矩阵，并从稠密表示构建它：

  from torch_sparse import SparseTensor
  
  dense_adj = ...    #1
  sparse_adj = SparseTensor.from_dense(dense_adj)
  #1 稠密邻接矩阵
  

一般来说，使用稀疏张量的主要动机是节省内存，尤其是在处理大规模图或具有高比例零元素的矩阵时。但如果你的数据中零元素很少，稠密张量在内存访问和计算速度方面可能会提供一定的优势，因为稀疏张量的索引和访问开销可能会抵消空间节省。注意，从一种表示转换为另一种表示本身可能会消耗内存和处理能力。

示例

一个学区聘请GeoGrid来研究该学区荣誉学生在多个校园之间的关系。工作的一部分是建立一个社交网络，学生为节点，学生之间的关联为边。Barker博士正在研究学生的社交网络图，试图确定友谊形成的模式：

• 初步分析——Barker博士发现，在这个小社区中，几乎每个人都认识其他人。从原始数据来看，有1,000名学生（节点）和大约450,000个友谊（边）。Barker博士将现有边数与总的可能连接数进行比较：n(n-1)/2，其中n是节点数；这个值为499,500。由于现有的边数（450,000）接近总边数（499,500），他确定自己处理的是一个稠密图。

• 稠密表示——考虑到图的密度：

  • 邻接矩阵的大小为1,000 × 1,000。
  • 如果每个学生有一个包含10个属性（如成绩、俱乐部数量等）的特征向量，那么节点特征将存储在一个大小为1,000 × 10的张量中。
  • 由于图的稠密性质，邻接矩阵中有大量非零元素，Barker博士首先考虑使用稠密表示来提高计算效率。

• 内存考虑——然而，随着Barker博士研究的深入，他计划将更多学校的数据纳入数据集，预计图会变得更大，但不一定更稠密。他预见到，随着数据集规模的增加，稠密表示可能会变得更加消耗内存。

• 稀疏表示——为了应对这一潜在问题，他决定尝试稀疏表示。他使用torch_geometric.transforms.ToSparseTensor转换函数将当前的稠密图数据集转换为稀疏张量表示。

• 结果——转换后，他发现使用稀疏表示节省了大量内存，足以选择这种表示，特别是考虑到他未来的计划。尽管计算时间略有增加，但内存节省使得稀疏格式更适合他日益增长的数据集。

7.6 GNN算法的选择

选择合适的GNN算法对于确保机器学习任务的可扩展性和效率至关重要，尤其是在处理大规模图和有限的计算资源时。除了预测性能和任务适配性之外，考虑到可扩展性，选择GNN算法的两种方式是通过分析时间和空间复杂度以及评估一些关键指标。

7.6.1 时间和空间复杂度

我们通过使用“大O符号”来衡量时间和空间复杂度，这是一种数学简写，用来解释随着输入大小变化，函数的增长或下降速度。它像是函数或算法的速度计，告诉你当输入变得非常大或趋向某一特定值时，它们会如何表现。在机器学习工程和开发中，衡量算法效率时特别有用。

注：有关大O符号的更全面解释，请参见Goodrich等人[9]。此外，任何关于算法的入门书籍都应该涵盖这个主题。

我们还在附录中讨论了与图和图算法相关的时间和空间复杂度，但这里列出一些时间复杂度的大O符号例子，按升序排列：

• 常数时间复杂度，O(1) ——这是最佳情况，无论输入大小如何，算法总是需要相同的时间。一个例子是通过索引访问数组元素。
• 线性时间复杂度，O(n) ——算法的运行时间与输入大小成线性关系。一个例子是在数组中查找特定的值。
• 对数时间复杂度，O(log n) ——运行时间与输入大小呈对数关系。具有这种时间复杂度的算法效率非常高。一个例子是二分查找。
• 平方时间复杂度，O(n²) ——算法的运行时间与输入大小的平方成正比。一个例子是冒泡排序。

当你了解了如何评估大O时，你可以利用GNN算法作者提供的信息来进行评估。通常，在算法的发表中，作者会提供算法的步骤，可以用来进行大O分析。此外，作者也经常提供自己的复杂度分析。

现在我们已经介绍了大O的好处，接下来列出它的一些注意事项。进行GNN算法的独立或比较复杂度分析可能会很具挑战性，原因包括以下几点：

• 多样化的操作——GNN算法涉及各种操作，如矩阵乘法、非线性变换和池化。每个操作具有不同的复杂度，这使得很难提供单一的度量。此外，并非所有GNN都使用相同的操作，因此将它们并排比较可能有限。在文献中，通常在比较GNN时，比较的是一个主要操作，而不是整个算法。
• 实现细节——GNN算法的实际实现（如使用特定的库、硬件优化或并行计算策略）也会影响其复杂度。

例如，表7.3比较了Bronstein等人[10]中GCN与GraphSAGE的复杂度。这个比较具体看了一个操作（前向传播中的卷积-like操作）在某种输入图（稀疏图）上的表现。具体来说，Bronstein等人比较了操作Y = ReLU(A × W)的时间和空间复杂度。分解后，这个操作由两个主要阶段组成：

• 矩阵乘法（A × W） ——这意味着我们将矩阵A（可以是我们的输入数据）与矩阵X（算法试图优化的权重或参数）相乘，再与矩阵W相乘。矩阵乘法是一种变换数据的方式。
• 激活（ReLU） ——修正线性单元（ReLU）是一种激活函数，用于将非线性引入模型。实际上，ReLU会将矩阵乘法的结果，对于每个元素，如果值小于0，它将其设为0。如果大于0，ReLU保持不变。

表7.3 两种图算法的可扩展性比较：GCN与GraphSAGE

算法	时间复杂度	空间复杂度	内存/每个周期时间/收敛速度	备注
GCN	O( Lnd² )	O( Lnd + Ld² )	内存：差	每个周期时间：好
			优点：光谱卷积：高效且适用于大规模图	泛化能力强：适用于各种图相关问题
			缺点：由于需要存储整个邻接矩阵和节点特征，内存和时间复杂度较高
GraphSAGE	O( Lbd² k L )	O( bk L )	内存：好	每个周期时间：差
			优点：通过使用邻域采样和小批量处理解决了GCN的可扩展性问题	缺点：当采样节点多次出现在邻域中时可能会引入冗余计算
			每个批次保持O( bk L )节点在内存中，但损失仅对其中b个节点计算
n = 图中节点数	d = 节点特征表示的维度	L = 消息传递迭代次数或层数	k = 每次跳跃采样的邻居数	b = 小批量中的节点数

从这个比较中可以得出一个结论：虽然GCN的复杂度依赖于输入图中的所有节点数，GraphSAGE的复杂度与节点数无关，这在时间和空间性能上有了很大的改善。GraphSAGE通过采用邻域采样和小批量处理实现了这一点。

示例

GeoGrid的任务是根据各种城市因素预测一个区域是否会发生开发。图中的节点代表地理区域，而边则可以代表到设施、道路网络或其他已开发区域的距离：

• 团队分析——虽然当前项目仅涉及一个大都市区，但GeoMap希望在未来逐步扩展系统，覆盖全国，包括一个拥有数百万地理节点和数十亿边的数据库。每个节点都有一个特征向量，可能包括土地价值、公共交通的接近度和分区法规等属性。
• 由于当前图的大小、扩展计划以及及时预测的需求，GeoGrid的数据科学团队必须谨慎选择适当的GNN架构。
• GCN——GCN易于解释，但其时间复杂度O(Lnd)可能会在图扩展时造成挑战。然而，通过使用PyG的小批量方法，团队可以管理图，而不需要存储整个邻接矩阵，使得GCN成为一个合理的候选方案。
• GraphSAGE——GraphSAGE提供O(Lbdk)的时间复杂度，因其内存效率和可扩展性而受到青睐。它允许调整小批量大小b和采样邻居数k，在性能调整上提供灵活性。
• GAT——图注意力网络（GAT）通过注意力机制提供了细致的见解，但其计算成本较高。尽管其大O复杂度可能与GCN类似，但注意力机制可能会增加额外的计算开销。

算法比较

虽然GCN看起来比GraphSAGE更简单，但它对节点数n的依赖可能会随着图的增长而成为问题。GraphSAGE因其对b和k的依赖而具备良好的可扩展性。GAT，尽管可能更准确，但由于其注意力机制，计算复杂度较高。

使用PyG进行小批量处理使得GCN更加易于管理。然而，团队也喜欢GraphSAGE，因为它固有的可扩展性优势。尽管GAT可能提供更高的准确性，但对于这个应用来说，它可能过于资源密集。

决策——经过彻底评估后，GeoGrid团队决定GraphSAGE提供了最平衡的方法，在计算效率和预测准确性之间进行了优化。

结论——他们计划稍后在受控环境中试用GAT，以评估其增加的计算需求是否确实带来了更准确的城市发展预测。他们将在生产前进行用户接受度测试，并设置明确的度量标准。

前面三节已经涵盖了在计划训练GNN时需要做出的基本选择，特别是在考虑规模问题时。接下来的五节将回顾可以解决规模问题的方法，包括深度学习优化、采样、分布式处理、使用远程后端和图粗化。

7.7 使用采样方法进行批处理

在本节中，我们探讨如何将大规模数据拆分为通过采样方法选择的批次。我们将首先进行一般性说明，然后详细介绍PyG包中的一些实现方法。最后，我们将通过一个GeoGrid案例，突出使用这些方法的实际选择和影响。

采样：文献与实现

在文献中，有很多关于各种采样技术（通常分类为节点采样、层采样和图采样）在GNN算法中设计的讨论，但在本节中，我们将重点关注PyG包中的采样实现。许多这些技术来源于文献，但它们的目的是将采样推广，以支持各种GNN算法和训练操作。在本节中，我们使用这些采样实现来支持小批量处理。

GCN是一个很好的示例。尽管标准形式的GCN模型确实不涉及采样，PyG的NeighborSampler函数仍然可以与GCNConv层一起使用。这是可能的，因为NeighborSampler本质上是一个数据加载器，它从更大的图中返回一个子图的批次。

在这个上下文中，子图被用来近似完整的图卷积操作。显而易见的优势是，我们可以处理大图，否则它们可能会超出算法或我们机器的内存。缺点是，使用NeighborSampler的GCNConv的准确性可能不如全批次训练，因为存在这种近似。

7.7.1 两个概念：小批量处理和采样

小批量处理和采样是两种不同的方法，通常可以将它们合并为一个函数。小批量处理（由PyG中的加载器完成）是将一个大数据集拆分为节点或边的子集，以便通过训练过程。但如何确定包含在较小组中的节点或边的子集呢？采样是我们用来选择这些子集的具体机制。这些子集可以是连接的子图，但不一定必须是。通过这种方式进行小批量处理将减轻内存负担。在一个周期中，我们可以在不将整个图存储在内存中的情况下，每次只存储它的一小部分。

通过采样进行的小批量处理可能有一些缺点。一个问题是关键信息的丢失。例如，如果我们考虑消息传递过程，每个节点及其邻域对于更新节点信息至关重要。采样可能会错过重要的节点，从而影响模型的性能。这就像在消息传递框架中遗漏了至关重要的消息。此外，采样过程可能引入偏差，影响模型的泛化能力。这相当于在消息传递框架中进行有偏的聚合操作。

在PyG中实现小批量处理

小批量处理方法可以在加载器和采样器模块中找到。大多数这些方法将采样方法与批处理功能结合在一起，将采样数据提供给模型训练过程。有一些基础类允许你编写自定义采样器（如baseloader、basesampler），以及具有预设采样机制的加载器[11, 12]。

选择合适的采样器

选择理想的采样方法可能不是一件简单的事，这取决于图的特性和训练目标。不同的采样器将导致不同的周期时间和收敛时间。没有通用规则可以确定最佳的采样器；最好通过对你的数据进行有限的实验，看看哪个方法最有效。实现采样增加了GNN架构的复杂性，就像消息传递需要精心编排的聚合和更新步骤一样。

7.7.2 快速了解PyG中的重要采样器

如我们所见，GNN通过聚合局部邻域来工作。然而，对于非常大的图，考虑在聚合操作中包含所有节点或边可能是不可行的，因此通常使用采样器。以下是一些常见的采样器，它们也由PyG库默认支持：

• NeighborLoader——理想用于捕捉局部邻域动态，常用于社交网络分析。
• ImbalancedSampler——为不平衡数据集设计，如欺诈检测场景。
• GraphSAINT变种——旨在最小化梯度噪声，使其适合大规模训练[9]。
• ShaDowKHopSampler——用于采样更大的邻域，捕捉更广泛的结构信息。
• DynamicBatchSampler——设计用于按邻居数量分组节点，优化批量计算一致性。
• LinkNeighborLoader——一个采样器，使用类似于NeighborLoader的方法采样边。

注：本概述并不全面，功能可能会根据使用的PyG版本有所不同。有关详细信息，请参阅PyG官方文档（mng.bz/DMBa）。

让我们看一个使用NeighborLoader加载器的代码片段。完整代码可以在GitHub仓库中找到，以下是一些片段。代码运行了一个使用采样器的GNN训练循环。对于每个批次，它将节点特征、标签和邻接信息传输到设备（通常是GPU）。然后，它清除先前的梯度，执行前向和反向传播计算损失，并相应地更新模型参数。要在代码中添加使用NeighborSampler的小批量采样，可以按照以下步骤操作：

1. 导入所需的模块：

   from torch_geometric.loader import NeighborLoader
   

2. 定义小批量大小和每个节点采样的层数：

   batch_size = 128    #1
   num_neighbors = 2    #2
   #1 设置所需的小批量大小
   #2 设置每个节点的采样层数
   

3. 创建NeighborLoader实例，在小批量训练过程中对邻域进行采样：

   loader = NeighborLoader(data, input_nodes = train_mask, batch_size=batch_size,
      num_neighbors=num_neighbors)
   

   这里，data是输入图，input_nodes包含训练节点的索引，num_neighbors指定每层采样的邻居数。

4. 修改你的训练循环，以使用采样器迭代小批量，如下所示：

代码7.1 使用NeighborSampler的训练循环

for batch_size, n_id, adjs in sampler:    #1
  x = data.x[n_id].to(device)        #2
  y = data.y[n_id].squeeze(1).to(device)   #3
  adjs = [adj.to(device) for adj in adjs]  #4

  optimizer.zero_grad()              #5
  out = model(x, adjs)             #6
  loss = F.nll_loss(out, y)         #7
  loss.backward()               #8
  optimizer.step()     #9

• #1 启动训练循环，通过NeighborSampler迭代批次。batch_size是批次大小，n_id包含节点ID，adjs存储采样子图的邻接信息。
• #2 获取当前批次中节点的特征（x），并将其传送到目标设备（通常是GPU）。这类似于在消息传递框架中获取嵌入。
• #3 获取当前批次中节点的标签（y），移除任何单一维度，并将其传送到设备。
• #4 将采样子图的邻接信息传送到设备。
• #5 将所有优化变量的梯度设置为零。这是反向传播期间正确计算梯度的必要步骤。
• #6 通过GNN模型进行前向传播计算预测。模型接收节点特征和邻接信息作为输入。
• #7 使用负对数似然损失计算模型输出与真实标签之间的损失。
• #8 进行反向传播，计算基于损失的梯度。
• #9 基于计算出的梯度更新模型参数。

在本节结束时，我们将看一个GeoGrid团队需要在三个批处理器中做出选择的案例。

示例

让我们回到GeoGrid，这是一家领先的地图公司。一个团队正在开发一个基于图的美国公路系统表示，交叉口作为节点，道路段作为边。这项项目的规模庞大，带来了计算和内存方面的挑战。

经过彻底调查后，团队挑选了三种主要的小批量处理技术，下面我们将评估每种方法的权衡：

• GraphSAINTSampler：它的噪声降低能力非常有利，能够提供更准确的梯度估计，并且具有良好的可扩展性，非常适合像美国公路网络这样的大规模系统。然而，它的实现可能较为复杂，且存在对高度连接的节点过度表示的风险。
• NeighborSampler：它在内存使用方面非常高效，关注关键的道路段，并强调局部邻域连接，能够提供有关重要交叉口的洞察。然而，它可能会忽略一些来自较少交通的路线的重要数据，且可能会对密集连接的节点产生偏向。
• ShaDowKHopSampler：它有效地采样k跳子图，捕捉更大的邻域，并且其深度可调，能够适应不同公路系统的复杂性。然而，某些k值可能会使其在计算上变得较为昂贵，并且广泛的捕捉可能会引入过多且不立即相关的数据。

接下来，我们展示如何在实际应用中使用不同的采样器，以GeoGrid公司为案例研究：

决策——经过广泛的讨论，团队倾向于选择ShaDowKHopSampler。该方法能够捕捉更广泛的邻域，而不局限于直接的邻居，似乎非常适合美国公路系统的多样化复杂性。他们认为，通过实验确定合适的k值，可以在深度和计算效率之间找到平衡。

为了避免潜在的信息过载并确保相关性，GeoGrid计划将结果与现实世界的交通数据进行对比，确保所采样的图保持实际性和准确性。

结论——GeoGrid选择采用ShaDowKHopSampler，这一决策源于他们深入分析了自己的需求与每种技术的优缺点。通过将采样方法与现实世界数据相结合，他们旨在在图表示的粒度和相关性之间找到平衡。

现在我们已经理解了小批量处理，我们可以研究两种与采样密切相关的技术：并行处理和使用远程后端。

7.8 并行与分布式处理

小批量处理方法非常适合与接下来的两种方法结合使用：并行处理和使用远程后端，因为这些方法在数据被拆分时效果最佳。并行处理是通过将计算任务分布到多个计算节点或多台机器上来训练机器学习模型的方法。在本节中，我们将重点讨论如何在单台机器上的多个GPU之间分配模型训练任务[13-17]。我们将使用PyTorch的DistributedDataParallel来实现这一目的。

数据并行与分布式数据并行

在PyTorch中，您将遇到两种主要的并行化神经网络模型的方法：DataParallel和DistributedDataParallel。每种方法都有其优点和局限性，这对于做出明智的决策至关重要。

• DataParallel：适用于单台机器上的多GPU设置，但存在一些缺点，例如每次前向传播时模型复制会产生额外的计算开销。随着模型和数据的扩展，这些限制会变得更加明显。
• DistributedDataParallel：跨多个机器和GPU扩展。通过为GPU之间的通信分配专用的CUDA缓冲区，通常带来较少的开销，因此其性能优于DataParallel。这使得它非常适合大规模数据和复杂的模型。

DataParallel和DistributedDataParallel都提供了并行化模型的途径。理解它们各自的优缺点，使您能够选择最适合您的机器学习挑战的技术。考虑到其在可扩展性和效率方面的优势，特别是在处理复杂或大规模项目时，我们选择DistributedDataParallel作为我们的首选模型并行化方法。

7.8.1 使用分布式数据并行

简而言之，分布式数据并行（DDP）是一种同时在多个图形卡（GPU）上训练机器学习模型的方法。其思想是将数据和模型分布到不同的GPU上，进行计算，然后将结果合并。为了使其工作，您首先需要设置一个进程组，这只是组织您使用的GPU的一种方式。与其他一些方法不同，DDP不会自动拆分数据，您需要自己进行这部分的操作。

当您准备好训练时，DDP通过同步所有GPU上的模型更新来帮助完成。通过共享梯度，所有GPU都可以接收这些更新，它们共同帮助改进同一个模型，尽管它们处理的是不同的数据块。

这种方法特别快速且高效，尤其是与单GPU训练或使用更简单的并行方法相比。然而，您需要注意一些技术细节，例如在使用多台机器时确保正确加载和保存模型。训练的基本步骤如下：

• 模型实例化：初始化用于训练的GNN模型。
• 分布式模型设置：使用PyTorch的DistributedDataParallel将模型包装起来，准备进行分布式训练。
• 训练循环：实现一个训练循环，包括前向传播、计算损失、反向传播和更新模型参数。
• 进程同步：使用PyTorch的分布式通信包同步所有进程，确保所有进程完成训练后再继续执行下一步。可以使用dist.barrier()在进入下一个周期前等待所有进程完成。所有周期完成后，销毁进程组。
• 入口点保护：使用if __name__ == '__main__':来指定数据集并启动分布式训练。这确保只有当脚本直接运行时，训练代码才会执行，而不会在作为模块导入时执行。

使用分布式处理需要仔细处理同步点，以确保模型正确训练。您还必须确保您的机器或集群有足够的资源来处理并行计算。

torch.distributed支持多种用于分布式计算的后端。最推荐的两个后端如下：

• NVIDIA Collective Communications Library (NCCL) ：NVIDIA的NCCL用于基于GPU的分布式训练。它提供了优化的原语，用于集体通信。
• Gloo：Gloo是一个集体通信库，由Facebook开发，提供各种操作，如广播、全规约等。这个库用于CPU训练。

7.8.2 DDP代码示例

以下是使用PyTorch进行分布式训练的示例。为简便起见，我们使用修改后的国家标准与技术研究所（MNIST）数据集训练一个简单的神经网络。使用GCN和Amazon Products数据集的示例可以在GitHub仓库中找到。在那个例子中，我们使用Kaggle笔记本，而不是Google Colab，这里有一个双GPU系统。GCN示例中的另一个不同之处在于，我们使用了NeighborLoader数据加载器，它使用了NeighborSampler采样器。

让我们来逐步解释这段代码。GCN版本本质上遵循相同的逻辑。

为分布式训练做准备

脚本导入必要的模块，如torch、torch.distributed等。它使用dist.init_process_group初始化DDP环境。它设置使用NCCL进行通信，并指定一个本地地址和端口（tcp://localhost:23456）用于同步。

准备模型和数据

代码定义了一个简单的Flatten层，这是神经网络的一部分，用于重塑输入。数据转换和加载步骤通过PyTorch的DataLoader和torchvision数据集进行设置。加载的数据是MNIST。

训练函数

train是负责训练模型的函数。它通过数据的小批量进行迭代，执行前向和反向传播，更新模型参数。

主函数

在main()函数内，每个进程（在这个例子中代表一个GPU）设置它的随机种子和设备（基于进程的rank设置CUDA设备）。神经网络模型被定义为一个顺序模型，先是Flatten层，然后是Linear层。然后它被包装在DistributedDataParallel中。损失函数（CrossEntropyLoss）和优化器（SGD）被定义。

多进程生成

最后，脚本使用mp.spawn函数启动分布式训练。它在world_size数量的进程上运行main()（基本上是两个GPU）。每个进程将在其数据子集上训练模型。

运行训练

每个进程使用其数据子集训练模型，但梯度会在所有进程（GPU）之间同步，以确保处理器更新相同的全局模型。这个过程在图7.3中进行了总结。

以下代码示例使用DistributedDataParallel模块训练一个神经网络。

代码7.2 使用DDP进行训练

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel    #1
from torch.utils.data import DataLoader    #2
from torchvision import datasets, transforms

class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)

def train(model, trainloader,
          criterion,
          optimizer,
          device):   #3
    model.train()
    for batch_idx, (data, target) in enumerate(trainloader):
        print(f'Process {device}, Batch {batch_idx}')
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

def main(rank, world_size):    #4
    filepath = '~/.pytorch/MNIST_data/'
    dist.init_process_group(   #5
        backend='nccl', 
        init_method='tcp://localhost:23456', 
        rank=rank,
        world_size=world_size    #6
    )

    torch.manual_seed(0)   #7
    device = torch.device(f'cuda:{rank}')    #8

    transform = transforms.Compose(
        [transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))]
    )

    trainset = datasets.MNIST(filepath,                #9
                               download=True,          #9
                               train=True,             #9
                               transform=transform)    #9
    train_loader = DataLoader(trainset,           #10
                               batch_size=64,    #10
                               shuffle=True,     #10
                               num_workers=2)    #10

    model = nn.Sequential(Flatten(), nn.Linear(784, 10)).to(device)
    model = DistributedDataParallel(model, device_ids=[rank])   #11

    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

    train(model, train_loader, criterion, optimizer, device)    #12

• #1 导入DistributedDataParallel类，用于分布式训练
• #2 导入DataLoader工具，用于数据加载
• #3 定义主要的训练函数
• #4 定义分布式训练设置的主函数
• #5 初始化分布式进程组
• #6 指定参与进程的总数
• #7 设置随机种子以确保结果可复现
• #8 根据进程的rank设置设备
• #9 加载并转换MNIST数据集
• #10 创建训练数据的小批量加载器
• #11 将模型包装在DistributedDataParallel中，准备分布式训练
• #12 调用训练函数启动训练过程

我们结束这一节，通过GeoGrid的一个实际示例，来展示他们如何做出决定。

示例

GeoGrid有机会为一个政府项目提交概念验证，旨在使用GNN进行复杂的环境建模。赢得这个合同可能会使他们成为这一领域的领导者，但他们面临着激烈的竞争。政府设定了一个紧迫的截止日期来审查概念验证演示，这使得GeoGrid的情况变得紧张，因为他们仍处于开发的早期阶段。

在一次团队会议上，焦点转向了一个关键的技术决策和一个重要的困境：是否在多个GPU上使用DDP训练。首席数据科学家看到了DDP加速训练时间的潜力，认为这可能为政府项目提供一个令人印象深刻的效率演示。

另一方面，团队中的一位经验丰富的工程师对DDP有一些担忧。尽管它有很多优点，但DDP可能会引入问题，例如在GPU之间同步梯度时的计算开销。另一层复杂性来自其他团队成员，他们指出他们的专业GNN算法尚未在DDP下进行过测试。他们对数据如何分配到GPU之间以及可能出现的不平衡和低效表示担忧。其他担忧则集中在开发和测试代码所需的时间上。

团队认真权衡了这些因素。快速按时制作演示是可取的。然而，将DDP应用于他们特定的GNN模型的复杂性和未知因素可能会导致意外的延迟和成本，甚至可能错过提交的截止日期。

团队进一步考虑了模型开发的迭代特性。在概念验证阶段，快速的性能优化迭代至关重要。将DDP引入其中可能会使调试变得更加复杂，并延长开发周期：

决策——最终，团队选择了一种审慎的方法。他们决定进行为期一周的可行性研究，严格评估使用DDP对他们的GNN架构的影响。这将使他们能够基于经验数据做出明智的决策，这些数据将跟踪收敛时间和每个周期的平均时间。IT部门将被咨询，确保有足够的计算资源专门用于这个关键的研究。

结论——决定是否推出GNN通常高度依赖于数据、时间表和计算需求。可行性研究是决策过程中的一个重要组成部分，尤其是在识别计算需求时。

在下一节中，我们将探讨另一种基于采样的技术：从远程存储系统直接抽取数据进行训练。

7.9 使用远程存储进行训练

本书中一种显著的数据管道方法是从数据存储系统中获取数据，并通过转换该数据使其可用于GNN平台进行预处理。这个预处理后的数据在训练过程中被存储在内存中。

然而，当数据太大以至于无法完全加载到内存时，一种方法是将预处理过程集成到训练过程中。我们不再将整个数据集预处理后加载到内存再进行训练，而是直接从初始数据存储系统中进行采样并小批量处理数据进行训练。通过GNN平台与数据源之间的接口，我们可以处理每个从数据源直接提取的批次[18]。在PyG中，这称为远程后端，它设计为与所使用的特定后端无关[19-22]。

其好处是，数据集的大小现在仅受限于数据库的容量。其权衡如下：

• 我们需要做一些工作来设置远程后端，本节将详细介绍。
• 从远程后端提取数据会引入I/O延迟。
• 集成远程后端增加了训练设置的复杂性。基本上，更多的事情可能出错，调试的项目也会更多。

在PyG中，远程后端通过从图的两个方面存储和采样来实现：使用GraphStore存储结构信息（即边），使用FeatureStore存储节点特征（在编写时，边特征尚未得到支持）。对于存储图结构，PyG团队建议使用图数据库作为后端，如Neo4J、TigerGraph、Kùzu和ArangoDB。同样，对于节点特征，PyG团队建议使用键值数据库，如Memcached、LevelDB和RocksDB。实现远程后端的关键要素如下：

• 远程数据源：存储图结构和节点特征的数据库。这个选择可能只是你当前用来存储图形的数据库系统。
• GraphStore对象：torch_geometric.data.GraphStore对象存储图的边索引，支持节点采样。自定义类的核心组件必须包括与数据库的连接，以及CRUD（创建、读取、更新、删除）功能，包括put_edge_index()、get_edge_index()和remove_edge_index()。
• FeatureStore对象：torch_geometric.data.FeatureStore管理图节点的特征。节点特征的大小被认为是图学习应用中的一个主要存储问题。像GraphStore一样，自定义实现包括连接远程数据库和CRUD功能。
• 采样器：与GraphStore链接的图采样器使用采样算法通过torch_geometric.sampler.BaseSampler接口从输入节点生成子图。PyG的默认采样器拉取边索引，将它们转换为压缩稀疏列（CSC）格式，并使用内存中的采样例程。自定义采样器可以通过实现BaseSampler类中的sample_from_nodes()和sample_from_edges()方法，使用GraphStore的专用方法。这涉及节点级别和边级别的采样。
• 数据加载器：数据加载器与之前章节中介绍的类似。不同之处在于，数据加载器使用创建的GraphStore、FeatureStore和采样器对象，而不是通常的PyG数据对象。以下是PyG文档中的一个示例。

代码7.3 使用远程后端的加载器对象

loader = NodeLoader(
    data=(feature_store, graph_store),
    node_sampler=node_sampler,
    batch_size=20,
    input_nodes='paper',
)

for batch in loader:
    <training loop>

虽然可以开发自定义类和功能，但建议使用数据库供应商提供的工具。目前，KuzuDB和ArangoDB为PyG的远程后端提供了实现[14, 18-20, 23]。我们结束本节时，再次展示GeoGrid的一个小案例。

7.9.1 示例

GeoGrid有一个如此庞大的图，无法完全装入现有硬件的内存中。他们希望使用GNN来分析这个大图，预测交通拥堵、路线受欢迎程度等特征。但如何在一个无法完全装入内存的图上训练GNN呢？以下是一些处理大规模GNN的具体示例：

• 采用PyG的远程后端：GeoGrid使用了PyG的远程后端功能，这与公司需要处理大规模图的需求完美契合。他们使用Neo4J作为图数据库来存储图结构，使用RocksDB来存储节点特征，例如位置类型、历史交通数据等。
• 远程数据源：GeoGrid选择了Neo4J和RocksDB作为他们的数据存储系统。首先的任务是编写脚本，将庞大的图数据加载到这些数据库中。这涉及数据验证，确保加载的数据是正确且一致的。
• GraphStore对象：GeoGrid的开发团队花费了大量时间来实现GraphStore对象。他们需要建立与Neo4J数据库的安全可靠连接。一旦连接建立，他们就实现了CRUD操作。
• FeatureStore对象：同样，实现RocksDB的FeatureStore对象也不简单。主要挑战是处理不同大小和类型的节点特征，这需要经过彻底的测试以确保效率和正确性。
• 采样器：开发自定义采样策略是一个单独的项目。采样器需要既有效又高效，经过多次迭代，最终达到了性能标准。
• 数据加载器：NodeLoader是最后拼图，将之前的所有元素组合成一个连贯的训练管道。开发团队必须确保NodeLoader针对速度进行了优化，以最小化I/O延迟。

测试和调试

与所有软件开发、机器学习或AI项目一样，测试是工作流中的关键部分。以下是处理项目时典型的测试和质量保证（QA）步骤：

• 单元测试：每个组件都经过严格的单元测试。这对于尽早发现错误至关重要，确保系统的每个部分在独立运行时按预期工作。
• 集成测试：单元测试后，团队进行了集成测试，运行整个管道，从加载一批数据到将其通过GNN模型运行。他们发现了几个瓶颈和错误，尤其是在采样器和数据库连接部分，这些问题花了不少时间来调试和解决。
• I/O延迟：公司遇到的一个重大问题是从Neo4J和RocksDB拉取数据时的I/O延迟。GeoGrid优化了查询，并使用了一些缓存机制来缓解这一问题。
• 调试：在开发和测试阶段，团队遇到了各种错误和问题，从数据不一致到采样过程中的意外行为。每个问题都需要仔细调试，增加了整体开发时间。

尽管面临这些挑战，GeoGrid成功地实现了一个可扩展的解决方案，用于在他们庞大的地理图上训练GNN。这个项目耗时且具有复杂性，但能够在内存不足的图上进行训练的可扩展性和能力是无法估量的好处，足以证明这一努力的价值。

7.10 图粗化

图粗化是一种用于减少图大小的技术，同时保留其基本特征。这项技术通过创建原始图的粗化版本来减少图的大小和复杂度。图粗化减少了节点和边的数量，使得图更加易于管理和分析。它通过聚合或合并节点和边来形成简化的原始图表示，同时尽量保留其结构和关系信息。

图粗化的一种方法是从输入图G（及其标签Y）开始，然后通过以下步骤生成粗化图G' [23]：

1. 对图G应用图粗化算法，生成一个标准化的划分矩阵（即节点簇集合）P。

2. 使用这个划分矩阵来完成以下操作：

   • 构建粗化图G'。
   • 计算G'的特征矩阵。
   • 计算G'的标签。
   • 使用粗化图进行训练，生成一个权重矩阵，并可在原始图上进行测试。

虽然我们可以通过减少顶点和边来使用图粗化技术来减小大图的大小，但它也有一些缺点。图粗化可能会导致信息丢失，因为原始图的关键信息可能被移除，进而使后续分析变得复杂。它还可能引入不准确性，因为无法完全表示原始图的结构。最后，没有一种通用的图粗化方法，因此可能会导致不同的结果并引入偏差。

在PyG中，图粗化涉及两个步骤：

• 聚类：这涉及将相似的节点聚集在一起，形成超节点。每个超节点代表原始图中的一组节点。聚类算法根据某些标准来确定哪些节点是相似的。在PyG中，有多种聚类算法可用，如graclus()和voxel_grid()。
• 池化：一旦形成了簇或超节点，池化就用于从原始图中创建一个粗化图。池化将每个簇中节点的信息合并为粗化图中的一个节点。在PyG中，max_pool()和avg_pool()是池化操作，它们输入来自第一步的簇。

如果重复使用，聚类和池化的组合允许我们创建一系列图，每一个比上一个更加简化，如图7.4所示。

7.10.1 示例

GeoGrid面临一个巨大的任务：分析美国道路系统的广泛图，以支持他们的雄心勃勃的交通管理解决方案。初始数据集包含50,000个节点和200,000条边，计算负担非常沉重。在初步探索时，GeoGrid考虑到计算负载，图粗化看起来是一个诱人的策略。但他们有很多顾虑。最初的担忧包括信息丢失以及考虑到标签保持和方法偏差的复杂性，可能引入的不准确性。

GeoGrid决定谨慎推进，使用Graclus算法和max_pool进行图的粗化试验。试验结果证实了公司的担忧。图的大小被显著减少，但以失去高流量区域的细节为代价。为聚类节点生成的新标签没有最佳地反映原始图，影响了机器学习模型的表现。

鉴于试验结果不理想，GeoGrid探索了其他优化方法。GeoGrid的突破性想法是建立一个多层次的分析框架，如下所示：

• 国家级——广泛的高层次层次，每个节点代表一个州或主要区域
• 州级——一个中间层，代表城市或县
• 城市级——最细粒度的层次，关注个别交叉口和道路段

团队推测，在中间层应用图粗化可能会缓解一些初期的顾虑。州级成为公司粗化的目标，这在计算效率和数据完整性之间提供了平衡。在这一新方法的指导下，GeoGrid重新评估了图粗化的缺点：

• 细粒度信息丢失——尽管仍然是一个问题，但由于粗化应用于中间层，因此损害似乎最小化，保留了城市级的细节。
• 引入不准确性——GeoGrid推测，其他层次可以作为补偿机制，弥补州级引入的不准确性。
• 标签保持——在州级进行粗化似乎风险较小，因为它们可以参考国家级和城市级进行修正。

他们继续使用相同的Graclus算法和max_pool技术对州级进行粗化。随后的评估发现，细粒度损失在此特定层次上是可以接受的，而引入的不准确性大部分被城市级和国家级平衡。

尽管公司最初回避了图粗化，但GeoGrid找到了将其有意义地融入更复杂的多层次系统中的方法。妥协使GeoGrid能够节省计算资源，而不会严重损害模型的准确性。然而，他们仍然保持谨慎，并致力于持续的研究，以全面了解涉及的权衡。

表7.5 使用图粗化的权衡分析，结合GeoGrid案例的见解

类别	见解	GeoGrid的应用案例
计算效率	适用于具有有限计算资源的实时处理	在州级加速了分析，减少了计算负担
简化分析	对于初步理解或宏观决策非常有用，适用于高层次概述	国家级层次提供了广泛的视图，作为低层次详细分析的基础
可扩展性	可以处理较大的图，避免计算上不可行的情况	多层次方法可以进一步扩展，包含额外的层次，如有需要
灵活性	可以应用于图的特定层次或部分，而非整个图	只对州级层次应用粗化，减少了一些缺点，同时仍然获得计算效益
细粒度信息丢失	不适用于需要精确、详细数据的任务	最初避免在交叉口层次进行粗化，因为丢失了关键的细节
潜在不准确性	需要从更详细的层次或附加数据中进行验证，以减轻不准确性	城市级和国家级作为对粗化后的州级的校验
标签保持挑战	需要额外的步骤来生成或映射新标签，这可能引入错误	在对中间层进行粗化时，发现标签更容易进行对接
方法偏差	选择粗化算法可能影响结果并引入偏差	识别为一个持续研究的领域，以更好地理解其影响

随着我们结束这一节，显而易见，扩展数据集的能力对于从事GNN工作的人员至关重要。处理大规模数据问题需要精心的策略，本节提供了应对这些挑战的多种方法的详细概述。从选择理想的处理器，到关于稀疏与密集表示的决策，从批量处理策略到分布式计算——扩展优化的选项是多种多样的。

在继续前进时，我们仓库中提供的代码可以作为一个有用的基准，确保这里提到的方法不仅仅是高层次的想法，而是可操作的计划。

在GNN的广阔领域中航行，需要战略眼光和动手执行的结合。无论你的数据大小或复杂度如何，诀窍在于规划、优化和迭代。让我们的见解成为你的指南针，帮助你自信地应对各种挑战，无论它们的规模如何。

总结

• 当训练非常大数据集时，时间和规模优化方法至关重要。我们可以通过图的顶点和边的原始数量、边和节点特征的大小，或者在处理和训练数据集时使用的算法的时间和空间复杂度来表征一个大图。

存在一些著名的技术来管理规模问题，这些技术可以单独使用或结合使用：

• 选择处理器及其配置

• 使用稀疏表示与密集表示的数据集

• 选择GNN算法

• 基于数据采样的批量训练

• 使用并行或分布式计算

• 使用远程后端

• 对图进行粗化

• 如何选择图数据的表示方式进行训练会影响性能。PyTorch Geometric (PyG) 提供了对稀疏和密集表示的支持。

• 训练算法的选择会影响训练的时间性能和内存的空间需求。使用Big O符号和基准测试关键指标可以帮助您选择最优的GNN架构。

• 节点或图批处理可以通过在训练中使用数据的部分而不是完整数据集来提高时间和空间复杂度。

• 并行处理，即将训练工作分配到一台机器的多个处理节点或多个机器的集群上，可以提高执行速度，但需要额外的设备设置和配置开销。

• 远程后端直接从外部数据源（图数据库和键/值存储）拉取数据，以进行训练中的小批量处理。这可以缓解内存问题，但需要额外的工作来设置和配置。

• 图粗化可以通过用更小的图代替原始图来减少内存需求。这个较小的版本是通过合并节点创建的。该方法的缺点是粗化图将偏离原始图的表示。图粗化是计算效率与数据忠实度之间的权衡。它最有效的应用是在谨慎使用并作为更大层次分析策略的一部分时。对中间层的应用可以缓解一些缺点。