graphstorm

GraphStorm 是一个面向企业用例的图机器学习（GML）框架。它通过为极大规模图（以十亿计的节点和边）提供可扩展的图机器学习（GML）模型训练和推理流程，简化了行业级图的 GML 模型开发、训练和部署。GraphStorm 提供了一系列内置的 GML 模型，用户无需编写任何代码即可通过单个命令训练 GML 模型。为了帮助开发最先进的模型，GraphStorm 提供了大量配置选项，用于定制模型实现和训练流程以提高模型性能。GraphStorm 还提供了一个编程接口，可以以分布式方式训练任何自定义 GML 模型。用户提供自己的模型实现，并使用 GraphStorm 训练流程进行扩展。

GraphStorm 架构

入门

安装

GraphStorm 兼容 Python 3.7+。它需要 PyTorch 1.13+、DGL 1.0 和 transformers 4.3.0+。

GraphStorm 可以通过 pip 安装，并可用于在独立模式下训练 GNN 模型。要在分布式环境中运行 GraphStorm，我们建议用户使用 Docker 容器来减少环境设置工作。设置 GraphStorm 运行环境的指南可以在这里找到，完整的分布式训练设置说明可以在这里找到。

使用 OGB 数据集运行 GraphStorm

注意：我们假设用户已按照使用 pip 包设置 GraphStorm 的说明设置了 GraphStorm 独立环境。并且用户已将 GraphStorm 源代码 git 克隆到 /graphstorm/ 文件夹中以使用一些补充工具。

OGB arxiv 图上的节点分类 首先，使用以下命令下载 OGB arxiv 数据并将其处理成用于节点分类任务的 DGL 图。

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

其次，使用以下命令将这个 arxiv 图划分为 GraphStorm 可以用作输入的分布式图。

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

GraphStorm 训练依赖 ssh 来启动训练任务。GraphStorm 独立模式使用 22 端口的 ssh 服务。

第三，运行以下命令在划分的 arxiv 图上训练 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

OGB MAG 图上的链接预测 首先，使用以下命令下载 OGB MAG 数据并将其处理成用于链接预测任务的 DGL 图。预测的边类型是 "author,writes,paper"。该命令还将此类型的 80% 的边用于训练和验证（默认 10%），其余 20% 用于测试。

python /graphstorm/tools/gen_mag_dataset.py --savepath /tmp/ogbn-mag-lp/ --edge-pct 0.8

其次，使用以下命令将 MAG 图划分为分布式格式。

python /graphstorm/tools/partition_graph_lp.py --dataset ogbn-mag \
                                               --filepath /tmp/ogbn-mag-lp/ \
                                               --num-parts 1 \
                                               --num-trainers-per-machine 4 \
                                               --target-etypes author,writes,paper \
                                               --output /tmp/ogbn_mag_lp_train_val_1p_4t

第三，运行以下命令在划分的 MAG 图上训练 RGCN 模型来执行链接预测。

python -m graphstorm.run.gs_link_prediction \
       --workspace /tmp/ogbn-mag-lp/ \
       --num-trainers 1 \
       --num-servers 1 \
       --num-samplers 0 \
       --part-config /tmp/ogbn_mag_lp_train_val_1p_4t/ogbn-mag.json \
       --ssh-port 22 \
       --cf /graphstorm/training_scripts/gsgnn_lp/mag_lp.yaml \
       --node-feat-name paper:feat \
       --save-model-path /tmp/ogbn-mag/models \
       --save-perf-results-path /tmp/ogbn-mag/models

要了解 GraphStorm 的全部功能，请参考我们的文档和教程。

引用

如果您在科学出版物中使用 GraphStorm，我们将感谢您引用以下论文：

@article{zheng2024graphstorm,
  title={GraphStorm: all-in-one graph machine learning framework for industry applications},
  author={Zheng, Da and Song, Xiang and Zhu, Qi and Zhang, Jian and Vasiloudis, Theodore and Ma, Runjie and Zhang, Houyu and Wang, Zichen and Adeshina, Soji and Nisa, Israt and others},
  journal={arXiv preprint arXiv:2406.06022},
  year={2024}
}