从信用卡交易、社交网络到推荐系统，再到生物学中的运输网络和蛋白质间的相互作用，图形是建模和分析复杂连接的首选数据结构。图形神经网络 (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 |
这些数字利用 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）。
由于可以将 ogbn-product 数据集加载到 GPU 显存中，速度甚至可以加快 115 倍(图 2)。
对于大型数据集(例如 ogn-papers100M)，必须开启 UVA 模式。
图 3 显示了节点分类任务中 ogn-papers100M 的每轮训练时间(以秒为单位)。它具有 1.11 亿个节点和 32 亿个边缘。
您可以在 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 的数据采样和训练性能。
要开始使用，请下载以下资源：
