GraphSAINT：基于图采样的归纳学习方法

Hanqing Zeng*, Hongkuan Zhou*, Ajitesh Srivastava, Rajgopal Kannan, Viktor Prasanna

联系方式

Hanqing Zeng (zengh@usc.edu), Hongkuan Zhou (hongkuaz@usc.edu)

欢迎随时报告问题或提出建议！

概述

GraphSAINT 是一个通用且灵活的框架，用于在大规模图上训练 GNN。GraphSAINT 提出了一种新颖的 mini-batch 方法，专门针对具有复杂关系的数据（即图）进行了优化。传统的 GNN 训练方式是：1) 在完整的训练图上构建 GNN；2) 对于每个 mini-batch，选择输出层中的一些节点作为根节点，然后从根节点反向追踪层间连接，直到到达输入层；3) 基于根节点上的损失进行前向和反向传播。而 GraphSAINT 的训练方式是：1) 对于每个 mini-batch，从完整的训练图中采样一个小子图；2) 在这个小子图上构建一个完整的 GNN，GNN 层内不再进行采样；3) 基于子图节点上的损失进行前向和反向传播。

GraphSAINT 采用的是“图采样”式的训练，而其他方法则是“层采样”式的训练。为什么改变采样的视角如此重要呢？GraphSAINT 能够实现以下几点：

准确性：我们通过简单而有效的归一化操作，消除了图采样带来的偏差。此外，由于任何采样过程都会因为丢弃邻居而导致信息丢失，我们提出了轻量级的图采样器，能够根据拓扑特征保留重要的邻居。实际上，图采样也可以被理解为一种数据增强或训练正则化手段（例如，边采样可以看作是 DropEdge 的 mini-batch 版本）。

效率：对于许多基于层采样的方法来说，“邻居爆炸”是一个令人头疼的问题，而 GraphSAINT 则凭借其图采样的理念提供了一个干净的解决方案。由于每个 GNN 层都是完整的且未经过采样，无论网络深度如何，邻居的数量始终保持不变。因此，每批次的计算成本会随着 GNN 深度的增加从指数级降低到线性级。

灵活性：GraphSAINT 中 mini-batch 子图上的层传播与在完整图上的传播几乎完全相同。因此，大多数为完整图设计的 GNN 架构都可以无缝地使用 GraphSAINT 进行训练。相比之下，一些层采样算法仅支持有限数量的 GNN 架构。以 JK-Net 为例，跳跃式知识连接要求浅层的节点样本必须是深层节点样本的超集——而 FastGCN 和 AS-GCN 的 mini-batch 格式并不满足这一条件。

可扩展性：GraphSAINT 在以下三个方面实现了可扩展性：1) 图规模：我们的子图大小无需随训练图规模成比例增长。因此，即使面对百万节点的图，子图仍然可以轻松容纳在 GPU 内存中；2) 模型规模：通过解决“邻居爆炸”问题，训练成本与 GNN 的宽度和深度呈线性关系；3) 并行资源数量：图采样可以通过简单的任务并行化实现高度可扩展。此外，解决“邻居爆炸”问题还意味着通信开销的大幅减少，这在分布式环境中尤为重要（参见我们发表在 IEEE/IPDPS '19 上的文章，或 硬件加速器开发 相关工作）。

【新闻】：请查看我们最新的工作，它将子图采样推广到了训练和推理两个阶段：shaDow-GNN（NeurIPS'21）！

关于本仓库

本仓库包含了我们两篇论文的源代码（ICLR '20 和 IEEE/IPDPS '19，详见【引用与致谢】部分）。

./graphsaint 目录包含 ICLR '20 论文中提出的 mini-batch 训练算法的 Python 实现。我们提供了两种实现版本，分别基于 Tensorflow 和 PyTorch。这两个版本遵循相同的算法。需要注意的是，我们论文中的所有实验均基于 Tensorflow 版本。而在公开图基准上的新实验则采用了 PyTorch 版本。

./ipdps19_cpp 目录包含了 IEEE/IPDPS '19 论文中描述的并行训练技术的 C++ 实现（详见 ./ipdps19_cpp/README.md）。本仓库的其余部分则专注于 ICLR '20 中的 GraphSAINT。

本仓库支持的 GNN 架构如下：

GNN 架构	Tensorflow	PyTorch	C++
GraphSAGE	:heavy_check_mark:	:heavy_check_mark:	:heavy_check_mark:
GAT	:heavy_check_mark:	:heavy_check_mark:
JK-Net	:heavy_check_mark:
GaAN		:heavy_check_mark:
MixHop	:heavy_check_mark:	:heavy_check_mark:

本仓库支持的图采样器如下：

采样器	Tensorflow	PyTorch	C++
节点	:heavy_check_mark:	:heavy_check_mark:
边	:heavy_check_mark:	:heavy_check_mark:
RW	:heavy_check_mark:	:heavy_check_mark:
MRW	:heavy_check_mark:	:heavy_check_mark:	:heavy_check_mark:
全图	:heavy_check_mark:	:heavy_check_mark:

其中：

• RW：随机游走采样器
• MRW：多维随机游走采样器
• 全图：始终返回完整的训练图，用作基线，实际上并未进行真正的采样。

您可以轻松添加自己的采样器和 GNN 层。详情请参阅【自定义】部分。

结果

新增：我们正在 Open Graph Benchmark 上测试 GraphSAINT。目前，我们已有 ogbn-products 图的数据结果。需要注意的是，排行榜上使用其他方法训练的 ogbn-products 准确率大多是在直推式设置下得出的。而我们的结果是在归纳式学习设置下得到的（归纳式学习难度更高）。

ICLR '20 中的所有结果都可以通过运行 ./train_config/ 中的配置文件来复现。例如，./train_config/table2/*.yml 存储了我们论文中表 2 的所有配置文件。./train_config/explore/*,yml 则存储了用于更深 GNN 和各种 GNN 架构（GAT、JK 等）的配置文件。此外，与 OGB 相关的结果是通过 ./train_config/open_graph_benchmark/*.yml 中的配置文件训练得到的。

测试集 F1-mic 分数汇总如下。

采样器	深度	GNN	PPI	PPI (large)	Flickr	Reddit	Yelp	Amazon	ogbn-products
Node	2	SAGE	0.960		0.507	0.962	0.641	0.782
Edge	2	SAGE	0.981		0.510	0.966	0.653	0.807
RW	2	SAGE	0.981	0.941	0.511	0.966	0.653	0.815
MRW	2	SAGE	0.980		0.510	0.964	0.652	0.809
RW	5	SAGE		0.995
Edge	4	JK				0.970
RW	2	GAT			0.510	0.967	0.652	0.815
RW	2	GaAN			0.508	0.968	0.651
RW	2	MixHop				0.967
Edge	3	GAT							0.8027

依赖项

• python >= 3.6.8
• tensorflow >=1.12.0 / pytorch >= 1.1.0
• cython >=0.29.2
• numpy >= 1.14.3
• scipy >= 1.1.0
• scikit-learn >= 0.19.1
• pyyaml >= 3.12
• g++ >= 5.4.0
• openmp >= 4.0

数据集

我们论文中使用的所有数据集均可下载：

• PPI
• PPI-large（PPI 的更大版本）
• Reddit
• Flickr
• Yelp
• Amazon
• ogbn-products
• ...（更多即将添加）

这些数据集可通过 Google Drive 链接 下载（或者通过 百度网盘链接（提取码：f1ao））。请将下载的文件夹重命名为根目录下的 data 文件夹。目录结构应如下所示：

GraphSAINT/
│   README.md
│   run_graphsaint.sh
│   ...
│
└───graphsaint/
│   │   globals.py
│   │   cython_sampler.pyx
│   │   ...
│   │
│   └───tensorflow_version/
│   │   │    train.py
│   │   │    model.py
│   │   │    ...
│   │
│   └───pytorch_version/
│       │    train.py
│       │    model.py
│       │    ...
│
└───data/
│   └───ppi/
│   │   │    adj_train.npz
│   │   │    adj_full.npz
│   │   │    ...
│   │
│   └───reddit/
│   │   │    ...
│   │
│   └───...
│

我们还提供了一个脚本，可以将我们的数据格式转换为 GraphSAGE 格式。要运行该脚本，

python convert.py <数据集名称>

例如，python convert.py ppi 会将 PPI 数据集转换为 GraphSAGE 格式，并保存到 ./data.ignore/ppi/。

新增：对于从 OGB 格式到 GraphSAINT 格式的转换，请使用 ./data/open_graph_benchmark/ogbn_converter.py 脚本。目前，该脚本可以处理 ogbn-products 和 ogbn-arxiv。

Cython 实现的并行图采样器

我们有一个需要编译的 Cython 模块，才能开始训练。请在根目录下运行以下命令来编译该模块：

python graphsaint/setup.py build_ext --inplace

训练配置

训练所需的超参数可以通过配置文件设置：./train_config/<name>.yml。

用于复现表 2 结果的配置文件已打包在 ./train_config/table2/ 中。

有关配置文件格式的详细说明，请参阅 ./train_config/README.md。

运行训练

首先，请编译 Cython 采样器（见上文）。

我们建议查看 ./graphsaint/globals.py 中定义的可用命令行参数（TensorFlow 和 PyTorch 版本共用）。通过正确设置这些标志，您可以最大化采样步骤中的 CPU 利用率（通过指定可用核心数量），选择日志文件的存放目录，并开启或关闭日志记录器（Tensorboard、Timeline 等）等。

注意：对于论文中比较的所有方法（GraphSAINT、GCN、GraphSAGE、FastGCN、S-GCN、AS-GCN、ClusterGCN），采样或聚类仅在训练过程中进行。为了获得验证集/测试集的准确率，我们会对整个图（包括训练、验证和测试节点）运行全批量 GNN，并仅针对验证/测试节点计算 F1 分数。详情请参阅 issue #11。

为简化实现，在验证/测试集评估时，我们使用完整的图邻接矩阵进行层传播。对于 Amazon 或 Yelp 数据集，这可能会导致某些 GPU 内存不足的问题。如果出现内存不足错误，请使用 --cpu_eval 标志强制在 CPU 上进行验证/测试集评估（而小批量训练仍将在 GPU 上进行）。其他标志说明如下。

要在 CPU 上运行代码

python -m graphsaint.<tensorflow/pytorch>_version.train --data_prefix ./data/<数据集名称> --train_config <train_config yml 路径> --gpu -1

要在 GPU 上运行代码

python -m graphsaint.<tensorflow/pytorch>_version.train --data_prefix ./data/<数据集名称> --train_config <train_config yml 路径> --gpu <GPU 编号>

例如，--gpu 0 表示在第一块 GPU 上运行。此外，还可以使用 --gpu <GPU 编号> --cpu_eval 来让 GPU 执行小批量训练，而 CPU 执行验证/测试评估。

我们还实现了双 GPU 训练以进一步加快运行速度。只需添加 --dualGPU 标志，并使用 --gpu 标志指定两块 GPU 即可。目前，此功能仅适用于支持显存池化且通过 NvLink 连接的 GPU。

新增：我们已准备了专门用于训练 OGB 图的脚本。相关脚本和说明请参阅 ./graphsaint/open_graph_benchmark/。

自定义

以下我们将介绍如何根据您的研究或产品需求自定义此代码库。

如何准备您自己的数据集？

假设您的完整图包含 N 个节点。每个节点有 C 个类别，以及长度为 F 的初始属性向量。如果您的训练/验证/测试划分比例为 a/b/c（即 a+b+c=1），则：

adj_full.npz：一个以 CSR 格式存储的稀疏矩阵，类型为 scipy.sparse.csr_matrix。其形状为 N×N。矩阵中的非零元素对应于完整图中的所有边。连接两节点的边是属于训练、验证还是测试节点并不重要。对于无权图，非零元素均为 1。

adj_train.npz：一个以 CSR 格式存储的稀疏矩阵，类型为 scipy.sparse.csr_matrix。其形状同样为 N×N。然而，矩阵中的非零元素仅对应于连接两个训练节点的边。图采样器只会从这个 adj_train 中选取节点和边，而不会从 adj_full 中选取。因此，在训练过程中既不会泄露属性信息，也不会泄露结构信息。另外请注意，只有 adj_train 中 aN 行和 aN 列包含非零元素。参见问题 #11。对于无权图，非零元素均为 1。

role.json：一个包含三个键的字典。键 'tr' 对应所有训练节点的索引列表，键 'va' 对应所有验证节点的索引列表，键 'te' 对应所有测试节点的索引列表。请注意，在原始数据中，节点可能具有字符串类型的 ID。您需要为节点重新分配数值 ID（0 至 N-1），以便能够通过索引访问邻接矩阵、特征矩阵和类别标签。

class_map.json：一个长度为 N 的字典。每个键是节点的索引，每个值要么是一个长度为 C 的二进制列表（用于多分类问题），要么是一个整数标量（0 至 C-1，用于单分类问题）。

feats.npy：一个形状为 N×F 的 numpy 数组。第 i 行对应节点 i 的属性向量。

如何添加您自己的采样器？

所有采样器都作为 GraphSampler 类的子类实现于 ./graphsaint/graph_samplers.py 文件中。您可以采用两种方式来实现自己的采样器子类：

1. 使用纯 Python 实现。重写父类的 par_sample 函数。我们在 ./graphsaint/graph_samplers.py 文件中的 NodeSamplingVanillaPython 类中提供了一个基本示例。
   • 优点：易于实现
   • 缺点：执行速度可能较慢。将纯 Python 函数并行化并非易事。
2. 使用 Cython 实现。您需要在 ./graphsaint/cython_sampler.pyx 文件中添加一个 Sampler 类的子类。在该子类中，您只需重写 __cinit__ 和 sample 函数。sample 函数定义了采样器的顺序行为。我们会自动通过同时启动多个采样器来实现任务级并行化。
   • 优点：可融入并行执行框架，执行速度接近 C++ 级别。
   • 缺点：编码难度较大。

如何支持您自己的 GNN 层？

在 ./graphsaint/<tensorflow 或 pytorch>_version/layers.py 文件中添加一层。此外，您还需要对 ./graphsaint/<tensorflow 或 pytorch>_version/models.py 文件中的 GraphSAINT 类的 __init__ 方法进行一些小的更新，以便模型能够根据 YAML 配置文件中的关键字查找正确的类。

引用与致谢

本研究得到 DARPA FA8750-17-C-0086 项目、NSF CCF-1919289 和 OAC-1911229 项目的资助。

我们感谢 Matthias Fey 在 PyTorch Geometric 库中提供的参考实现。

我们感谢 OGB 团队在大规模实验中使用 GraphSAINT。

• ICLR 2020：

@inproceedings{graphsaint-iclr20,
title={{GraphSAINT}: 基于图采样的归纳学习方法},
author={Hanqing Zeng 和 Hongkuan Zhou 和 Ajitesh Srivastava 和 Rajgopal Kannan 和 Viktor Prasanna},
booktitle={国际表征学习会议},
year={2020},
url={https://openreview.net/forum?id=BJe8pkHFwS}
}

• IEEE/IPDPS 2019：

@INPROCEEDINGS{graphsaint-ipdps19,
author={Hanqing Zeng 和 Hongkuan Zhou 和 Ajitesh Srivastava 和 Rajgopal Kannan 和 Viktor Prasanna},
booktitle={2019 年 IEEE 国际并行与分布式处理研讨会 (IPDPS)},
title={准确、高效且可扩展的图嵌入},
year={2019},
month={五月},
}

GraphSAINT 快速上手指南

GraphSAINT 是一个基于图采样（Graph Sampling）的通用框架，专为在大规模图上训练图神经网络（GNN）而设计。与传统基于“层采样”的方法不同，GraphSAINT 通过对每个 mini-batch 采样一个完整的子图来构建 GNN，有效解决了深层网络中的“邻居爆炸”问题，显著提升了训练效率、可扩展性和模型灵活性。

环境准备

系统要求

• 操作系统: Linux (推荐), macOS, Windows (需配置相应编译环境)
• Python: >= 3.6.8
• 编译器: g++ >= 5.4.0 (需支持 OpenMP >= 4.0)

前置依赖

请确保已安装以下核心库（建议使用虚拟环境）：

• 深度学习框架: TensorFlow >= 1.12.0 或 PyTorch >= 1.1.0
• 科学计算: NumPy >= 1.14.3, SciPy >= 1.1.0, Scikit-learn >= 0.19.1
• 其他工具: Cython >= 0.29.2, PyYAML >= 3.12

提示: 国内用户可使用清华或阿里镜像源加速安装，例如： pip install -i https://pypi.tuna.tsinghua.edu.cn/simple <package_name>

安装步骤

1. 克隆代码库

git clone https://github.com/GraphSAINT/GraphSAINT.git
cd GraphSAINT

2. 安装 Python 依赖

pip install -r requirements.txt
# 如果项目中没有 requirements.txt，请手动安装上述列出的依赖包

3. 编译 Cython 并行采样器

GraphSAINT 使用 Cython 加速图采样过程，使用前必须编译该模块。在项目根目录执行：

python graphsaint/setup.py build_ext --inplace

4. 准备数据集

下载论文中使用的数据集（PPI, Reddit, Flickr, Yelp, Amazon, ogbn-products 等）。

• 下载地址: Google Drive 或 百度网盘 (提取码: f1ao)
• 目录结构: 下载后解压，将文件夹重命名为 data 并放置于项目根目录下。最终结构应如下：

  GraphSAINT/
  ├── data/
  │   ├── ppi/
  │   ├── reddit/
  │   └── ...
  ├── graphsaint/
  └── ...
  

OGB 数据集转换: 若使用 Open Graph Benchmark (OGB) 格式的数据，请使用脚本转换： python ./data/open_graph_benchmark/ogbn_converter.py

基本使用

GraphSAINT 通过 YAML 配置文件管理超参数。配置文件位于 ./train_config/ 目录下。

1. 选择配置文件

• 复现论文 Table 2 结果：使用 ./train_config/table2/*.yml
• 探索深层网络或不同架构（如 GAT, JK-Net）：使用 ./train_config/explore/*.yml
• OGB 基准测试：使用 ./train_config/open_graph_benchmark/*.yml

2. 运行训练

在 GPU 上运行

指定数据集路径、配置文件路径及 GPU 编号：

python -m graphsaint.pytorch_version.train --data_prefix ./data/reddit --train_config ./train_config/table2/reddit_edge.yml --gpu 0

(注：将 pytorch_version 替换为 tensorflow_version 可切换框架)

在 CPU 上运行

将 --gpu 参数设为 -1：

python -m graphsaint.pytorch_version.train --data_prefix ./data/reddit --train_config ./train_config/table2/reddit_edge.yml --gpu -1

混合模式（GPU 训练 + CPU 评估）

若在验证/测试阶段遇到显存不足（OOM），可强制在 CPU 上进行评估：

python -m graphsaint.pytorch_version.train --data_prefix ./data/yelp --train_config ./train_config/table2/yelp_rw.yml --gpu 0 --cpu_eval

3. 自定义训练

• 修改超参数: 直接编辑对应的 .yml 配置文件。
• 添加新采样器或模型: 参考 graphsaint 目录下的现有实现，轻松扩展自定义的 GNN 层或图采样算法。