GraphStorm:面向企业级大规模图机器学习的开源框架

GraphStorm:赋能企业级大规模图机器学习

在当今数据驱动的时代,图数据结构因其强大的表达能力和建模复杂关系的能力而备受关注。然而,随着图规模的不断扩大,如何高效地在大规模图上进行机器学习成为了一个重要挑战。为了应对这一挑战,AWS Labs推出了GraphStorm - 一个专为企业级大规模图机器学习设计的开源框架。

GraphStorm的核心特性

GraphStorm的设计宗旨是简化在十亿级节点和边的大规模图上开发、训练和部署图机器学习(GML)模型的过程。它具有以下核心特性:

可扩展性: GraphStorm提供了可扩展的训练和推理管道,能够处理企业级的大规模图数据。
内置模型: 框架内置了多种GML模型,用户只需一个命令即可训练模型,无需编写代码。
灵活配置: 提供了丰富的配置选项,可以自定义模型实现和训练管道,以提升模型性能。
分布式训练: 支持自定义GML模型的分布式训练,用户可以利用GraphStorm的训练管道实现模型的可扩展性。
多任务支持: 支持节点分类、链接预测等多种图学习任务。

GraphStorm的架构设计

GraphStorm architecture

GraphStorm的架构设计充分考虑了大规模图处理的需求。它包括以下主要组件:

图构建模块: 负责将原始数据转换为分布式图格式。
模型训练模块: 包含多种内置GNN模型和可自定义的训练管道。
分布式训练引擎: 基于DGL和PyTorch,实现高效的分布式图学习。
推理模块: 支持模型在大规模图上的高效推理。
配置管理: 提供灵活的配置选项,以适应不同的应用场景。

快速上手GraphStorm

要开始使用GraphStorm,首先需要安装框架。GraphStorm兼容Python 3.7+,依赖PyTorch 1.13+、DGL 1.0和transformers 4.3.0+。用户可以通过pip安装GraphStorm:

pip install graphstorm

对于分布式环境,推荐使用Docker容器来简化环境配置。

以下是在OGB arxiv数据集上进行节点分类的示例:

下载并处理数据:

python /graphstorm/tools/gen_ogb_dataset.py --savepath /tmp/ogbn-arxiv-nc/ --retain-original-features true

对图进行分区:

python /graphstorm/tools/partition_graph.py --dataset ogbn-arxiv \
                                            --filepath /tmp/ogbn-arxiv-nc/ \
                                            --num-parts 1 \
                                            --num-trainers-per-machine 4 \
                                            --output /tmp/ogbn_arxiv_nc_train_val_1p_4t

训练RGCN模型:

python -m graphstorm.run.gs_node_classification \
       --workspace /tmp/ogbn-arxiv-nc \
       --num-trainers 1 \
       --part-config /tmp/ogbn_arxiv_nc_train_val_1p_4t/ogbn-arxiv.json \
       --ssh-port 22 \
       --cf /graphstorm/training_scripts/gsgnn_np/arxiv_nc.yaml \
       --save-perf-results-path /tmp/ogbn-arxiv-nc/models