异构图转换器 (HGT)

替代参考 Deep Graph Library (DGL) 实现

异构图转换器是一种图神经网络架构,可以处理大规模异构和动态图。

您可以查看我们的 WWW 2020 论文"异构图转换器"了解更多详情。

此 HGT 实现基于 Pytorch Geometric API

概述

本项目中最重要的文件如下:

conv.py: 我们模型的核心,实现了类似转换器的异构图卷积层。
model.py: 不同模型组件的封装。
data.py: 数据接口和用法。
- class Graph: 异构图的数据结构。在 Graph.node_feature 中以 pandas.DataFrame 存储特征;在 Graph.edge_list 中以字典存储邻接矩阵。
- def sample_subgraph: 异构图的采样算法。每次迭代对每种类型采样固定数量的节点。所有采样的节点都在已采样节点的区域内,采样概率为相对度的平方。
train_*.py: 特定下游任务的训练和验证脚本。
- def *_sample: 给定任务的采样函数。记得屏蔽图中已存在的链接以避免信息泄露。
- def prepare_data: 使用多个进程并行进行采样,可以与模型训练无缝协调。

设置

此实现基于 pytorch_geometric。要运行代码,您需要以下依赖项:

Pytorch 1.3.0
pytorch_geometric 1.3.2
- torch-cluster==1.4.5
- torch-scatter==1.3.2
- torch-sparse==0.4.3
gensim
sklearn
tqdm
dill
pandas

您可以简单地运行 pip install -r requirements.txt 来安装所有必要的包。

OAG 数据集

我们当前的实验在开放学术图谱 (OAG) 上进行。为了便于使用,我们将整个数据集分割并预处理成不同粒度:所有 CS 论文 (8.1G), 所有 ML 论文 (1.9G), 所有 NN 论文 (0.6G),跨度从 1900 年到 2020 年。您可以通过这个链接下载预处理的图。

如果您想直接从原始数据处理,可以通过这个链接下载。下载后,运行 preprocess_OAG.py 来提取特征并将它们存储在我们的数据结构中。

您也可以使用我们的代码处理其他异构图,只要将它们加载到 data.py 中的 class Graph 数据结构中。参考 preprocess_OAG.py 以获得演示。

用法

执行以下脚本,使用 HGT 训练论文-领域 (L2) 分类任务:

python3 train_paper_field.py --data_dir 数据集路径 --model_dir 保存模型路径 --conv_name hgt

进行其他两个任务的方法类似。这个脚本有一些关键选项:

conv_name: 选择相应的模型进行训练。默认我们使用 HGT。
--sample_depth 和 --sample_width: 采样图的深度和宽度。如果模型超出 GPU 内存,可以考虑减少它们的数量;如果想训练更深的 GNN 模型,考虑增加这些数字。
--n_pool: 并行进行采样的进程数。如果有大内存的机器,可以考虑增加这个数字以减少批次准备时间。
--repeat: 重复使用采样批次进行训练的次数。如果训练时间远小于采样时间,可以考虑增加这个数字。

其他可选超参数的详细信息可以在 train_*.py 中找到。

引用

如果您在应用中使用我们的代码,请考虑引用以下论文。

@inproceedings{hgt,
  author    = {Ziniu Hu and
               Yuxiao Dong and
               Kuansan Wang and
               Yizhou Sun},
  title     = {Heterogeneous Graph Transformer},
  booktitle = {{WWW} '20: The Web Conference 2020, Taipei, Taiwan, April 20-24, 2020},
  pages     = {2704--2710},
  publisher = {{ACM} / {IW3C2}},
  year      = {2020},
  url       = {https://doi.org/10.1145/3366423.3380027},
  doi       = {10.1145/3366423.3380027},
  timestamp = {Wed, 06 May 2020 12:56:16 +0200},
  biburl    = {https://dblp.org/rec/conf/www/HuDWS20.bib},
  bibsource = {dblp computer science bibliography, https://dblp.org}
}

pyHGT

异构图转换器 (HGT)

概述

设置

OAG 数据集

用法

引用