这篇文章最初发表在 NVIDIA 技术博客上。
从信用卡交易、社交网络到推荐系统，再到生物学中的运输网络和蛋白质间的相互作用，图形是建模和分析复杂连接的首选数据结构。图形神经网络 (GNN) 具有学习和推理图形结构化数据的能力，已在各个领域成为改变游戏规则的技术。
但是，发现这些图形中隐藏的模式和宝贵见解可能具有挑战性，尤其是在 GNN 的数据采样和端到端训练中。
为了弥补这一差距，NVIDIA 推出了GNN 框架，这是专为 DGL 和 PyG 设计的容器，具备以下特点：
- 用于数据采样的 GPU 加速
- GNN 训练和部署工具 (GNN 工具)
本文概述了 NVIDIA 加速 DGL 和 PyG 容器的优势，展示了客户在生产环境中使用这些容器的方式，并提供了性能指标。
在 NGC 目录中介绍 DGL 容器
深度图形库 (DGL) 是一款热门的开源库，它可以在现有的深度学习框架（如 PyTorch）上实现和训练图神经网络 (GNN)。
我们很高兴地宣布，DGL 现已通过其他 NVIDIA 库进行加速，并作为 容器 通过 NGC 目录 提供 ------ 这是一个 GPU 加速的 AI/ML、HPC 应用程序、SDK 和工具的中心。该目录使用户能够更快地访问性能优化的软件，并简化构建和部署 AI 解决方案的过程，帮助您更快地将产品推向市场。有关更多信息，请参阅 NGC 目录中用于 AI、数字孪生和 HPC 的数百个预训练模型（视频）。
23.09 版本的 DGL 容器提高了#DGL 用户的数据采样和训练性能。以下是此版本的主要功能。
用于数据加载程序采样的 GPU 加速
RAPIDS cuGraph 的 sampler 能够在几秒钟内处理数千亿个边，并同时为数千批世界上最大的 GNN 数据集计算样本。DGL 容器随附 cuGraph-DGL，这是 DGL 的加速扩展，使用户能够利用这一惊人的性能。
根据使用 8 个 V100 GPU 运行的基准测试，即使在中型数据集(约 10 亿个边缘)上，cuGraph 数据加载性能也至少比原生 DGL 快 2 -- 3 倍。cuGraph-DGL 采样通过在多个节点和多个 GPU 中分配图形，还能在此过程中节省内存，比线性扩展高达 1000 亿个边缘更好。
cuGraph 可以在 16 秒内对 1000 亿个边缘进行采样！ NVIDIA 专有库 cuGraph-ops 加速了 cuGraphSAGE、cuGraphGAT 和 cuGraphRGCN 等 GNN 运算符和模型，将模型前向时间缩短了一半。
GNN 训练和部署工具
GNN 工具是一个灵活的平台，可轻松训练和部署 GNN 模型。此工具基于热门的 Deep Graph Library (DGL) 和 PyTorch Geometric (PyG) 框架构建，使您能够为快速 GNN 实验构建端到端工作流。
它提供了一个完全模块化和可配置的工作流，可为自定义 GNN 用例实现快速迭代和实验。 NVIDIA 在我们的容器中包含示例 notebook，以便轻松进行实验。
多架构支持
NGC 中发布的 DGL 容器具有 x86 和 ARM64 版本，以支持新的 NVIDIA Grace Hopper GPU 这两个版本使用相同的容器标签。从基于 Arm 的 Linux 系统中提取容器时，您需要提取 ARM64 容器。
使用基于 ARM64 的 DGL 容器在 Grace Hopper 上训练 GNN 的示例
GNN 训练中的统一虚拟寻址 (UVA) 模式极大地受益于 NVIDIA Grace CPU 和 Grace Hopper 中 NVIDIA Hopper GPU 之间的连接。在 Grace Hopper 上，使用 ogn-papers100M 数据集训练同一 GraphSAGE 模型需要 1.9 秒/次，与使用 PCIe 连接的 H100*Intel CPU 训练相比，速度快了约 9 倍(表 1)。
| 系统 | GH200 | H100+Intel CPU | A100+AMD(DGX A100) | AMD Genoa(仅限 CPU) |
| (秒/次) | 1.9 | 16.92 | 24.8 | 107.11 |
表 1.GraphSage 模型每训练一次的训练时间与 OGBN 论文 100MGrace Hopper 上的数据集(UVA 模式)
这些数字利用 Grace Hopper 上的巨大页面来处理图形及其功能。它们在 512 GB Grace Hopper 节点上进行基准测试。该模型运行时批量大小为 4096，扇出 (3030) (在两层 GraphSAGE 模型中查看每个节点多达 30 个相邻节点)。这在 CUDA 12.1 上的 DGL 版本 1.1 上运行。
DGL 训练性能
训练 GNN 面临的挑战之一是数据加载过程。在某些情况下，例如使用 GraphSAGE 对 ogb-papers100M 数据集进行节点分类，数据加载过程占端到端训练时间的 90%以上。DGL 0.8 v 启用 UVA 模式，以实现图形特征的高效 GPU 加载，这提高了性能。
考虑使用 GraphSAGE 模型进行 ogbn-products 节点分类任务的数据集。该数据集包含约 240 万个节点和 6190 万条边。在 DGX-1 V100 GPU 上，采用 UVA 模式进行训练时，相比仅使用 CPU，它可以提供高达 20 倍的速度提升（见图 1）。
图 1.obgn-products 数据集 UVA 模式下的节点分类训练时间
由于可以将 ogbn-product 数据集加载到 GPU 显存中，速度甚至可以加快 115 倍(图 2)。
图 2.使用 GPU 显存训练 obgn-products 数据集的节点分类训练时间
对于大型数据集(例如 ogn-papers100M)，必须开启 UVA 模式。
图 3 显示了节点分类任务中 ogn-papers100M 的每轮训练时间(以秒为单位)。它具有 1.11 亿个节点和 32 亿个边缘。
图 3.obgn-papers 数据集的 UVA 模式下的节点分类训练时间
您可以在 NGC DGL 23.09 容器的/workspace/examples/多 GPU 目录。
PyTorch Geometric 容器
PyTorch Geometric (PyG) 是另一个热门的开源库，用于编写和训练各种应用的 GNN。我们将推出使用 cuGraph 等 NVIDIA 库加速的 PyG 容器。
根据使用 8 个 A100 GPU 运行的基准测试，在中型数据集(大约 10 亿个边缘)上，cuGraph 数据加载性能至少比原生 PyG 快 4 倍。
我们已经观察到一些客户从 PyG 容器中受益，并且我们计划利用 PyG 加速来与 NVIDIA BioNeMo 模型一起使用。
客户成功案例
以下介绍了不同的公司如何使用 NVIDIA 加速 DGL 和 PyG 容器来加速其工作流程。
适用于基于物理的 ML 的 GNN
NVIDIA Modulus 是一个开源框架，用于在 Python 中构建、训练和微调基于物理性质的 机器学习（ML） 模型。
随着 GNN 在计算流体动力学、分子动力学模拟和材料科学等学科中的应用兴趣与日俱增， NVIDIA Modulus 开始利用 DGL 和 cuGraphOps 库来支持 GNN.
NVIDIA Modulus 目前支持包括 MeshGraphNet、AeroGraphNet 和 GraphCast 在内的 GNN，用于基于网格的模拟和全球天气预报的模型。除了网络架构，Modulus 还提供了一系列工具，用于开发 天气预报、空气动力学模拟 以及 涡流脱落 的示例。
用于欺诈检测的 GNN
凭借数十年的经验，美国运通 在使用 AI 驱动的工具和模型监控及降低欺诈风险方面有着显著的成就。它们还能有效地识别信用卡行业中参与欺诈活动的个人。
American Express AI Labs 正在持续开展研究，通过实施基于图形的机器学习解决方案，更深入地了解欺诈者网络。
NGC 上发布的 DGL 容器使 AmEx 能够试验各种 GNN 架构，并大规模利用节点和边缘信息。这使得 AmEx 能够在多节点、多 GPU 环境中使用 NVIDIA 库处理数百万个节点和数十亿个边缘时实现计算效率。此外，用户友好且适应性强的库使他们能够轻松自定义组件，例如损失函数、采样技术等。
用于药物研发的 GNN
Astellas 作为一家以前瞻性药物研发方法而闻名的领先制药公司，正在利用 GNN 的功能来完成药物研发中的一系列关键任务。这些任务包括使用基于扩散模型的概率性深度学习模型来生成 3D 分子构象、进行特征提取和构建预测性 ML 模型，特别是在从头开始的蛋白质设计和工程方面。
因此，Astellas 利用 NVIDIA PyG 和 DGL 容器的计算效率，在药物研发活动中促进和放大基于 GNN 的 AI/ML 解决方案。由 PyG 和 DGL 提供支持的 AI 和基于 ML 的流程的集成使 Astellas 能够提高其内部能力，为研究人员和 AI 从业者提供开发和实施 AI 驱动的切割技术所需的专业知识。
作为一个用例，与基于仿真的传统 3D 分子构象方法相比，Astellas 科学家可以实现至少 50 倍的加速率。
Genentech，作为一家生物技术公司以及 Roche Group 的一部分，正在结合使用 GNN 和 NVIDIA PyG 容器，以加速其小分子预测训练。PyG 容器为 PyG 框架提供了稳定的基础，使得开发团队能够更加专注于开发工作，而不是搭建开发环境。
后续步骤
借助 NVIDIA 加速的 DGL 和 PyG 容器，您还可以显著提高 GNN 的数据采样和训练性能。
要开始使用，请下载以下资源：