我们非常高兴地宣布，GraphStorm 0.1 开源发布了。GraphStorm 是一个低代码企业级图形机器学习 (ML) 框架，可在几天内构建、训练和部署复杂企业级图形上的图形 ML 解决方案，而非几个月。使用 GraphStorm，您可以构建直接考虑数十亿个实体之间的关系或交互结构的解决方案，这些实体在大多数现实世界的数据中都是固有的，包括欺诈检测情形、推荐、社区检测和搜索/检索问题。
迄今为止，为复杂的企业级图形构建、训练和部署图形 ML 解决方案一直是非常困难的，这些图形往往具有数十亿个节点、数百亿条边和数十个属性——只需想想捕获 Amazon.com 产品、产品属性、客户和更多信息的图形即可。通过 GraphStorm，我们释放了亚马逊内部使用的将大规模图形 ML 解决方案推向生产的工具。GraphStorm 不需要您成为图形 ML 专家，并且在 GitHub 上以 Apache v2.0 许可证提供。有关 GraphStorm 的更多信息，请访问 GitHub 存储库。
在本文中，我们提供了 GraphStorm 的介绍、其架构以及如何使用的示例用例。
介绍 GraphStorm
图形算法和图形 ML 正在成为许多重要业务问题的最先进解决方案，例如预测交易风险、预测客户偏好、检测入侵、优化供应链、社交网络分析和交通预测。例如，亚马逊 GuardDuty，AWS 原生的威胁检测服务，使用数十亿条边的图形来提高其威胁情报的覆盖范围和准确性。这使得 GuardDuty 可以基于其与已知恶意域的关联将以前未见过的域分类为高度可能是恶意或良性的。通过使用图形神经网络 (GNN)，GuardDuty 能够增强其警报客户的能力。
然而，开发、发布和操作图形 ML 解决方案需要数月时间，并需要图形 ML 专业知识。首先，图形 ML 科学家必须使用 Deep Graph Library (DGL) 等框架为给定用例构建图形 ML 模型。由于企业应用程序中图形的大小和复杂性经常达到数十亿个节点、数百亿条边、不同的节点和边类型以及数百个节点和边属性，因此训练此类模型具有挑战性。企业图形可能需要千兆字节的存储器存储，这要求图形 ML 科学家构建复杂的训练流水线。最后，在模型被训练之后，它们必须被部署进行推断，这需要与训练流水线一样困难的推断流水线。
GraphStorm 0.1 是一个低代码企业图形 ML 框架，允许 ML 从业者轻松选择已被证明有效的预定义图形 ML 模型，在数十亿个节点的图形上运行分布式训练，并将模型部署到生产环境中。GraphStorm 提供了一系列内置的图形 ML 模型，例如用于具有异构图的企业应用程序的关系图卷积网络 (RGCN)、关系图注意力网络 (RGAT) 和异构图变换器 (HGT)，这允许具有很少图形 ML 专业知识的 ML 工程师尝试不同的模型解决方案，并快速选择正确的解决方案。端到端的分布式训练和推断流水线，可扩展到十亿级的企业图形，使训练、部署和运行推断变得容易。如果您是 GraphStorm 或图形 ML 的新手，您将从预定义的模型和流水线中受益。如果您是专家，则可以调整训练流水线和模型架构以获得最佳性能。GraphStorm 建立在 DGL 之上，这是一个广泛流行的用于开发 GNN 模型的框架，并以 Apache v2.0 许可证的形式提供开源代码。
“GraphStorm 的设计旨在帮助客户为行业应用程序尝试和操作图形 ML 方法，以加速图形 ML 的采用，”亚马逊 AI/ML 研究的高级主要科学家 George Karypis 表示。“自在亚马逊内部发布以来，GraphStorm 将构建基于图形 ML 的解决方案的工作量减少了多达五倍。”
“GraphStorm 使我们的团队能够在具有 2 十亿条边的图形上以自我监督的方式训练 GNN 嵌入，”亚马逊测量、广告技术和数据科学的首席应用科学家 Haining Yu 表示。“预训练的 GNN 嵌入在购物者活动预测任务上显示出比基于状态-of-the-art BERT 的基线的 24% 的改进；在其他广告应用中，它也超过了基准性能。”
“在 GraphStorm 之前，客户只能垂直扩展以处理 5 亿条边的图形，”亚马逊海王星和亚马逊时间流的 GM Brad Bebee 表示。“GraphStorm 使客户能够在拥有数十亿条边的大规模亚马逊海王星图形上扩展 GNN 模型训练。”
GraphStorm 技术架构
下图展示了 GraphStorm 的技术架构。
GraphStorm 基于 PyTorch 构建，可在单个 GPU、多个 GPU 和多个 GPU 机器上运行。它由三个层组成（在上述图中用黄色框标记）：
- 底层（Dist GraphEngine） – 底层提供了基本组件，以实现分布式图 ML，包括分布式图、分布式张量、分布式嵌入和分布式采样器。GraphStorm 提供了这些组件的高效实现，以将图 ML 训练扩展到数十亿个节点的图形。
- 中间层（GS training/inference pipeline） – 中间层提供了训练器、评估器和预测器，以简化内置模型和自定义模型的模型训练和推理。基本上，通过使用此层的 API，您可以专注于模型开发，而不必担心如何扩展模型训练。
- 顶层（GS general model zoo） – 顶层是一个模型仓库，为不同图形类型提供了流行的 GNN 和非 GNN 模型。截至本文撰写时，它为异构图提供了 RGCN、RGAT 和 HGT，为文本图提供了 BERTGNN。在未来，我们将添加对时间图模型（如用于时间图的 TGAT）以及知识图的 TransE 和 DistMult 的支持。
如何使用 GraphStorm
安装 GraphStorm 后，您只需要三个步骤就可以为您的应用程序构建和训练 GML 模型。
首先，您需要对数据进行预处理（可能包括自定义特征工程），并将其转换为 GraphStorm 所需的表格格式。对于每种节点类型，您需要定义一个表格，列出该类型的所有节点及其特征，并为每个节点提供一个唯一的 ID。对于每种边类型，您同样需要定义一个表格，其中每行包含该类型边的源节点和目标节点 ID（有关更多信息，请参见使用您自己的数据教程）。此外，您需要提供一个描述整个图结构的 JSON 文件。
其次，通过命令行界面（CLI），您使用 GraphStorm 的内置 construct_graph 组件进行一些 GraphStorm 特定的数据处理，从而实现高效的分布式训练和推理。
第三步，您需要在 YAML 文件中配置模型和训练（示例），并再次使用 CLI，调用五个内置组件之一（gs_node_classification、gs_node_regression、gs_edge_classification、gs_edge_regression、gs_link_prediction）作为训练管道来训练模型。此步骤将生成训练好的模型工件。要进行推理，您需要重复前两个步骤，将推理数据转换为使用相同的 GraphStorm 组件（construct_graph）构建的图形。
最后，您可以调用五个内置组件之一，与训练模型时使用的相同组件，作为推理管道来生成嵌入或预测结果。
整个流程也在以下图中说明。
在下一节中，我们提供一个使用示例。
在原始 OAG 数据上进行预测
在本文中，我们演示了如何使用 GraphStorm 轻松地对大型原始数据集进行图形 ML 训练和推理。开放学术图（OAG）包含五个实体（论文、作者、会议、机构和研究领域）。原始数据集以 JSON 文件的形式存储，大小超过 500 GB。
我们的任务是构建一个模型，以预测论文的研究领域。为了预测研究领域，您可以将其制定为多标签分类任务，但使用 one-hot 编码存储标签很困难，因为有数十万个领域。因此，您应该创建研究领域节点，并将此问题制定为链接预测任务，预测哪些研究领域节点应连接到论文节点。
要使用图形方法对该数据集进行建模，第一步是处理数据集并提取实体和边。您可以从JSON文件中提取五种类型的边来定义图形，如下图所示。您可以使用GraphStorm示例代码中的Jupyter笔记本来处理数据集，并为每种实体类型生成五个实体表和每种边类型生成五个边表。Jupyter笔记本还会对具有文本数据（如论文）的实体生成BERT嵌入。
在定义了实体和实体之间的边之后，您可以创建 mag_bert.json ，该文件定义了图形模式，并调用GraphStorm中内置的图形构建管道 construct_graph 来构建图形（见以下代码）。尽管GraphStorm图形构建管道在单个计算机上运行，但它支持多进程并行处理节点和边功能（ --num_processes ），并且可以将实体和边特征存储在外部存储器（ --ext-mem-workspace ）上以扩展到大型数据集。
python3 -m graphstorm.gconstruct.construct_graph \
--num-processes 16 \
--output-dir /data/oagv2.1/mag_bert_constructed \
--graph-name mag --num-partitions 4 \
--skip-nonexist-edges \
--ext-mem-workspace /mnt/raid0/tmp_oag \
--ext-mem-feat-size 16 --conf-file mag_bert.json要处理如此大的图形，您需要拥有具有大内存的CPU实例来构建图形。您可以使用Amazon Elastic Compute Cloud（Amazon EC2）r6id.32xlarge实例（128个vCPU和1 TB RAM）或r6a.48xlarge实例（192个vCPU和1.5 TB RAM）来构建OAG图形。
构建图形后，您可以使用 gs_link_prediction 在四个g5.48xlarge实例上训练链接预测模型。在使用内置模型时，您只需要调用一个命令行来启动分布式训练作业。请参见以下代码：
python3 -m graphstorm.run.gs_link_prediction \
--num-trainers 8 \
--part-config /data/oagv2.1/mag_bert_constructed/mag.json \
--ip-config ip_list.txt \
--cf ml_lp.yaml \
--num-epochs 1 \
--save-model-path /data/mag_lp_model在模型训练之后，模型工件保存在文件夹 /data/mag_lp_model 中。
现在，您可以运行链接预测推理来生成GNN嵌入并评估模型性能。GraphStorm提供多个内置评估指标来评估模型性能。例如，对于链接预测问题，GraphStorm自动输出平均倒数排名（MRR）指标。MRR是评估图形链接预测模型性能的有价值指标，因为它评估实际链接在预测链接中的排名高度。这捕捉到了预测的质量，确保我们的模型正确地优先考虑真实的连接，这是我们的目标。
您可以使用以下代码一次性运行推理。在这种情况下，该模型在构建的图形的测试集上达到了0.31的MRR。
python3 -m graphstorm.run.gs_link_prediction \
--inference --num_trainers 8 \
--part-config /data/oagv2.1/mag_bert_constructed/mag.json \
--ip-config ip_list.txt \
--cf ml_lp.yaml \
--num-epochs 3 \
--save-embed-path /data/mag_lp_model/emb \
--restore-model-path /data/mag_lp_model/epoch-0/请注意，推理管道从链接预测模型生成嵌入。要解决为任何给定论文找到研究领域的问题，只需对嵌入执行k最近邻搜索。
结论
GraphStorm是一种新的图形ML框架，可轻松在行业图形上构建、训练和部署图形ML模型。它解决了图形ML中的一些关键挑战，包括可扩展性和可用性。它提供了内置组件，可以从原始输入数据到模型训练和模型推理处理十亿级别的图形，并使多个亚马逊团队能够在各种应用中训练最先进的图形ML模型。请查看我们的GitHub存储库以获取更多信息。