图神经网络框架DGL v0.4.3 新版本发布

2021-02-01 10:18:22 浏览数 (1)

DGL团队昨天发布了 DGL v0.4.3 版本,这个版本无论是针对 DGL 系统的易用性还是性能,都做出了许多重要改进。主要亮点如下:

  • DGL 开始支持 TensorFlow 后端;
  • 发布 DGL-KE:高性能、易使用、可扩展的大规模知识图谱嵌入学习框架;
  • 发布 DGL-LifeSci:面向化学和生物领域的 GNN 算法库;
  • 发布新的图采样 API。

详细信息如下

TensorFlow 后端

DGL在这个版本开始正式拥抱 TensorFlow 社区,用户可以很容易地将 DGL 后端切换到 TensorFlow。如果你是第一次使用,在安装完 DGL 并使用 import dgl 导入后,就可以按照交互命令自动设置默认的后端。你也可以通过修改 ~/.dgl 目录下的 config.json 文件来修改后端引擎。对于用户来说,使用 DGL 并不需要关心当前后端采用的是哪一个引擎。如下代码展示了如何在 DGL 中使用 graph convolution layer.

代码语言:javascript复制
import tensorflow as tf
import dgl.nn as dglnn

# Random features for 10 nodes; each is of length 5.
x = tf.random.normal((10, 5))
# Random graph; 10 nodes and 20 edges.
g = dgl.rand_graph(10, 20)
# Pre-defined graph convolution module.
conv = dglnn.GraphConv(5, 3)
y = conv(g, x)  # Apply the graph convolution layer.

DGL团队实现并发布了 15 种常用的 Tensorflow GNN 模块(更多模块正在路上),所有的模块都可以只用一行代码调用。

  • GraphConv from the Graph Convolutional Networks paper.
  • GATConv from the Graph Attention Networks paper.
  • SAGEConv from the Inductive Representation Learning on Large Graphs paper (a.k.a. GraphSAGE).
  • GINConv from the How Powerful are Graph Neural Networks paper.
  • RelGraphConv from the Modeling Relational Data with Graph Convolutional Networks paper.
  • SGConv from the Simplifying Graph Convolutional Networks paper.
  • APPNPConv from the Predict then Propagate: Graph Neural Networks meet Personalized PageRank paper.
  • An edge_softmax function for computing softmax over the neighboring edges of each vertex.
  • Various pooling layers: SumPooling, AvgPooling, MaxPooling, SortPooling, WeightAndSum, and GlobalAttentionPooling.
  • A HeteroGraphConv module for applying GNN modules to heterogeneous graphs.

初步的 Benchmark 显示,相比于其他基于 Tensorflow 的 GNNs 实现,基于 DGL Tensorflow 后端的实现无论是训练速度还是内存消耗,都有很大的改进和提升。

DGL-KE:高性能、易使用、可扩展的大规模知识图谱嵌入学习框架

作为之前在 DGL 下的子项目,DGK-KE 目前正式作为独立软件包发布其 0.1 版本。主要亮点如下:

  • 支持当前主流的知识图谱嵌入表示模型,例如 TransE,DistMult,ComplEx,TransR,RotatE,RESCAL 等。
  • 用户无需编写代码,只需要将知识图谱数据作为输入,并调用 1 行命令即可训练。
  • DGL-KE是目前开源系统中唯一对多核 CPU,多卡 GPU,CPU-GPU 混合训练,以及分布式训练都支持的知识图谱嵌入表示框架。
  • 高效的 METIS 图分割算法,可以让分布式训练中的网络开销减小 90%,从而高效地训练大规模知识图谱数据。

用户可以通过 pip 直接安装 DGL-KE:

pip install dglke

以下命令展示了如何在 EC2 p3.16xlarge 上使用 8 台 GPU 训练 Freebase 知识图谱(86M nodes,338M 边)。

代码语言:javascript复制
dglke_train --model TransE_l2 --dataset Freebase --batch_size 1000 
--neg_sample_size 200 --hidden_dim 400 --gamma 10 --lr 0.1 --regularization_coef 1e-9 
-adv --gpu 0 1 2 3 4 5 6 7 --max_step 320000 --log_interval 10000 --async_update 
--rel_part --force_sync_interval 10000

DGL-KE 采用了许多创新的系统和算法优化. 我们的 benchmark 显示,对于 Freebase 这个大规模知识图谱数据集 (86M nodes,338M edges),DGL-KE 可以在 EC2 p3.16xlarge (8 GPU)上 100 分钟训练完模型。在 4 台 r5dn.24xlarge 上 30 分钟完成训练。这一速度比当前同类型的开源框架快 2-5 倍。

DGL-KE vs GraphVite on FB15k

DGL-KE vs Pytorch-Big-Graph on Freebase

DGL-KE 的 github 项目主页:

https://github.com/awslabs/dgl-ke

DGL-LifeSci:面向化学和生物领域的 GNN 算法库

作为之前在 DGL 项目下的 model zoo,DGL-LifeSci 现在也正式成为独立的软件包,其亮点主要如下:

  • 提供不同应用的训练脚本和预训练模型,包括:molecular property prediction, generative models, 以及 reaction prediction.
  • 相比于已有的其他工具实现具有 5.5 倍加速。
  • 良好的数据处理,模型构建,以及模型评测 pipeline

用户可以通过 pip 或者 conda 安装 DGL-LifeSci。

代码语言:javascript复制
pip install dgllife
conda install -c dglteam dgllife

Benchmark 结果如下:

DGL-LifeSci 的项目主页:

https://github.com/dmlc/dgl/tree/master/apps/life_sci

新的图采样API

图采样算法是在大图上训练 GNN 算法的关键。在这次的发布中,DGL团队重新实现了图采样 API,目标是为用户提供更加直观的编程体验以及更好的性能。新的 API 具有如下几点优势:

  • 支持一系列基于图采样的 GNN 模型,例如 PinSAGE,GraphSAGE,Graph Convolutional Matrix Completion (GCMC)等等。
  • 用户可以通过 python 代码自定义图采样算法。
  • 支持 Heterogeneous graphs。
  • 无缝兼容已有的 NN modules。
  • 通过多进程和多线程的配合来最大化提升性能。

如下代码展示了一个基本的 Neighbor sampler 实现:

代码语言:javascript复制
class NeighborSampler(object):
 def __init__(self, g, fanouts):
   self.g = g # The full graph structure
   self.fanouts = fanouts # fan-out of each layer


 def sample_blocks(self, seeds):
   # `seeds` are the set of nodes to build one sample from.
   blocks = []
   for fanout in self.fanouts:
     # For each seed node, sample ``fanout`` neighbors.
     frontier = dgl.sampling.sample_neighbors(g, seeds, fanout, replace=True)
     # Then we compact the frontier into a bipartite graph for message passing.
     block = dgl.to_block(frontier, seeds)
     # Obtain the seed nodes for next layer.
     seeds = block.srcdata[dgl.NID]


     blocks.insert(0, block)
   return blocks

尽管新的 API 还处于实验阶段,用户依然可以在很多示例中看到它的使用:

  • Train the GraphSAGE model by neighbor sampling and scale it to multiple GPUs .
  • Train the Relational GCN model on heterogeneous graphs by sampling for both node classification and link prediction.
  • Train the PinSAGE model by random walk sampling for item recommendation.
  • Train the GCMC model by sampling for MovieLens rating prediction.
  • Implement the variance reduction technique for neighbor sampling proposed by Chen et al.

DGL团队也在持续打磨新的 API,对应的文档和教程也会在随后发布。

其他改进

  • 在 dgl.nn 下的所有 GNN modules 现在都支持 Homogeneous Graph 和 Bipartite Graph。
  • DGLHeteroGraph 有了更快的 pickling/unpickling 实现。
  • 实现了新的加载和存储 DGLHeteroGraph 的 API。
  • BatchedDGLGraph 和 DGLSubGraph 类合并进去 DGLGraph 类。
  • 构建 DGLGraph 时不再需要 is_multigraph 参数。

0 人点赞