GraphStorm 03：可扩展的图形多任务学习，使用用户友好的 API 机器学习博客

GraphStorm 03：具有用户友好API的可扩展图形多任务学习

重要信息概述

GraphStorm 03 正式发布，为图形多任务学习提供原生支持。用户可以在一个训练循环中定义多个节点和边的训练目标，开发图形机器学习解决方案的速度大幅提升。此外，新版 GraphStorm 提供了简化的 API，用户只需编写 12 行代码即可实现自定义的节点分类训练循环。还发布了基于微软学术图MAG数据集的研究，展示了 GraphStorm 在图形处理上的性能和扩展性。

GraphStorm 是一个低代码企业级图形机器学习GML框架，旨在快速构建、训练和部署复杂企业级图形 ML 解决方案。在许多真实世界的数据中，亿万实体之间的关系结构嵌入其中，可应用于欺诈检测、推荐系统、社区检测和搜索/检索等场景。

支持图形多任务学习

许多企业应用涉及与不同节点和边相关联的图形数据。例如，零售组织希望在卖家和买家之间进行欺诈检测，而科学出版者则想寻找更多相关文献以引用。为更好地建模这些应用场景，客户要求支持图形多任务学习。

GraphStorm 03 支持六种常见的图形多任务学习，包括： 节点分类 节点回归 边分类 边回归 链接预测 节点特征重建

客户可以通过 YAML 配置文件指定训练目标。以下示例展示了科学出版者如何同时定义一个论文主题分类任务和一个论文引用的链接预测任务：

yamlversion 10gsf basic multitasklearning nodeclassification targetntype paper labelfield labelclass maskfields trainmaskclass valmaskclass testmaskclass numclasses 10 taskweight 10

     linkprediction        numnegativeedges 4                 numnegativeedgeseval 100          trainnegativesampler joint         trainetype             papercitingpaper            maskfields             trainmasklp                      valmasklp                        testmasklp                  taskweight 05

有关如何使用 GraphStorm 进行图形多任务学习的更多详尽信息，请参考 GraphStorm 的多任务学习文档。

新API自定义GraphStorm管道和组件

自2023年初 GraphStorm 发布以来，客户主要使用命令行界面CLI来快速构建和训练模型。然而，客户反馈希望有一个更易于定制训练和推断管道的接口。针对客户提出的需求，GraphStorm 03 引入重构后的图 ML 管道 API。使用新的 API，只需 12 行代码即可定义自定义的节点分类训练管道，以下示例展示了这一过程：

pythonimport graphstorm as gsgsinitialize()

acmdata = gsdataloadingGSgnnData(partconfig=/acmgs1p/acmjson)

traindataloader = gsdataloadingGSgnnNodeDataLoader(dataset=acmdata targetidx=acmdatagetnodetrainset(ntypes=[paper]) fanout=[20 20] batchsize=64)valdataloader = gsdataloadingGSgnnNodeDataLoader(dataset=acmdata targetidx=acmdatagetnodevalset(ntypes=[paper]) fanout=[100 100] batchsize=256 traintask=False)testdataloader = gsdataloadingGSgnnNodeDataLoader(dataset=acmdata targetidx=acmdatagetnodetestset(ntypes=[paper]) fanout=[100 100] batchsize=256 traintask=False)

model = RgcnNCModel(g=acmdatag numhidlayers=2 hidsize=128 numclasses=14)evaluator = gsevalGSgnnClassificationEvaluator(evalfrequency=100)

trainer = gstrainerGSgnnNodePredictionTrainer(model)trainersetupevaluator(evaluator)

trainerfit(traindataloader valdataloader testdataloader numepochs=5)

为帮助用户快速上手新 API，我们还发布了 新 Jupyter 笔记本示例，可在我们的 文档与教程页面 上找到。

LMGNN综合研究

许多企业应用的数据图中包含文本特征。例如，在零售搜索应用中，购物日志数据提供了文本丰富的产品描述、搜索查询和客户行为之间的关联。单靠基础大型语言模型LLMs并不适合处理这些数据，因为其学习的预训练数据分布和关系并不对应，而图机器学习GML则非常适合建模相关数据。但到目前为止，GML 从业者必须手动将 GML 模型与 LLM 结合起来，以获取最佳性能，特别是在底层图数据集较大时，这项手动工作既复杂又耗时。

在 GraphStorm 02 中，引入内置技术以高效结合语言模型LMs和图神经网络GNN模型，特别是在处理大量文本丰富图数据时。客户对 GraphStorm 的 LMGNN 技术在优化性能方面的应用提出了建议。为此，GraphStorm 03 发布了一个基于大型图数据集微软学术图MAG的 LMGNN 基准，涵盖节点分类和链接预测两个标准图 ML 任务。以下表格展示了该数据集的详细统计信息：

我们在 GraphStorm 中基准测试了两种主要的 LMGNN 方法，分别是广泛采用的预训练 BERTGNN 和 GraphStorm 开发者于 2022 年推出的微调 BERTGNN。通过预训练的 BERTGNN 方法，我们首先使用预训练的 BERT 模型计算节点文本特征的嵌入，然后训练 GNN 模型进行预测。通过微调 BERTGNN 方法，我们首先在图数据上微调 BERT 模型，并使用微调后的 BERT 模型计算的嵌入来训练 GNN 模型。实验中，我们使用 8 个 r524xlarge 实例进行数据处理，并使用 4 个 g548xlarge 实例进行模型训练和推断。微调的 BERTGNN 方法在链接预测MAG中的表现比预训练的 BERTGNN 方法提升了最高 40。

以下表格展示了这两种方法的模型性能及完整管道的整体计算时间，NC 代表节点分类，LP 代表链接预测。LM 时间成本表示计算 BERT 嵌入和微调 BERT 模型所花费的时间。

数据集任务数据处理时间目标预训练 BERT GNN微调 BERT GNNLM 时间成本每个周期时间指标LM 时间成本每个周期时间指标MAGNC553 minpaper subject206 min135 minLP198 mincite2195 minMrr 0487

我们还对大规模合成图进行了基准测试，以展示其可扩展性。我们生成了三个合成图，分别具有 10 亿、100 亿和 1000 亿边。对应的训练集大小分别为 800 万、8000 万和 8 亿。以下表格展示了图预处理、图分区和模型训练的计算时间。总体而言，GraphStorm 能够在数小时内对 1000 亿规模的图进行构建和训练！

有关更多基准测试详细信息和结果，请参见我们的 KDD 2024 论文。

总结

GraphStorm 03 在 Apache20 许可证下发布，旨在帮助您解决大规模图 ML 挑战，现在提供对多任务学习的原生支持以及用于自定义管道和 GraphStorm 组件的新 API。有关详细信息，请参考 GraphStorm GitHub 代码库 和 文档 开始您的旅程。

关于作者

Xiang Song 是 AWS AI 研究和教育AIRE部门的高级应用科学家，负责开发包括 GraphStorm、DGL 和 DGLKE 在内的深度学习框架。他主导了 Amazon Neptune ML 的开发，这是一项使用图神经网络的 Neptune 新功能。他目前正在主导GraphStorm的开发，这是一个旨在满足企业用例的开源图机器学习框架。他于 2014 年在复旦大学获得计算机系统与结构学博士学位。

Jian Zhang 是一名高级应用科学家，致力于利用机器学习技术帮助客户解决各种问题，如欺诈检测、装饰图像生成等。他成功开发了基于图的机器学习方案，特别是图神经网络，广泛应用于中国、美国和新加坡的客户。作为 AWS 图形能力的发言人，Zhang 曾多次公开演讲，介绍 GNN、深度图书馆DGL、Amazon Neptune 和其他 AWS 服务。

Florian Saupe 是 AWS AI/ML 研究的首席技术产品经理，支持图形机器学习组及大规模分布式训练、推理和故障恢复等系统组。在加入 AWS 之前，Florian 在博世负责自动驾驶的技术产品管理，曾是麦肯锡公司的战略顾问，并曾从事控制系统/机器人科学研究，并获得相关博士学位。

加载评论

• • 2026-01-27 12:03:03
  • 4