开源图神经网络框架DGL升级:GCMC训练时间从1天缩到1小时

乾明 编辑整理

量子位 报道 | 公众号 QbitAI

又一个AI框架迎来升级。

这次，是纽约大学、亚马逊联手推出图神经网络框架DGL。

不仅全面上线了对异构图的支持，复现并开源了相关异构图神经网络的代码，在GCMC、RCGN等业内知名的模型实现上，取得更好效果。

还发布了训练知识图谱嵌入（Knowledge Graph Embedding）专用包DGL-KE，在许多经典的图嵌入模型上都获得了性能提升。

目前，这个在2018年NeurIPS大会上发布的框架，已经获得了3000+标星，已经是图神经网络和图机器学习领域较受关注的工具。

异构图：神经网络的新方向

异构图是一个与同构图相对应的新概念。

传统同构图（Homogeneous Graph）数据中只存在一种节点和边，因此在构建图神经网络时所有节点共享同样的模型参数并且拥有同样维度的特征空间。

而异构图（Heterogeneous Graph）中可以存在不只一种节点和边，因此允许不同类型的节点拥有不同维度的特征或属性。

这一特点使得异构图的应用十分广泛，如果用图来描述我们和周围事物的关系就会发现所产生的图都是天然异构的，比如这个例子：

我今天看了电影《流浪地球》，那“我”作为观众和电影《流浪地球》之间就建立了“看了”这一关系。

异构图可以用来描述这种交互关系的集合。这个图分“观众”和“电影”两类节点，以及“看了”这一类边。

“我”作为观众，和电影所具有的属性一定是不同的，需要用不同的模型或者不同的特征维度来表达。

所以说，这张图就天然具有了异构性。而且，基于异构图训练的神经网络与传统方法相比，在一些场景中，也能够获得更好的效果与表现。

现在已经被应用到了知识图谱、推荐系统以及恶意账户识别等领域以及任务中。

其中比较有代表性的两个模型分别是用于节点分类和链接预测等任务的RGCN和用于产品推荐的GCMC。

但如何设计“既快又好”的深度神经网络？

这正是纽约大学、亚马逊联手推出图神经网络框架DGL的出发点。

新版DGL：异构图神经网络利器

这次更新放出的，是DGL 0.4版本，让整个框架变得更加实用。

不仅全面上线对于异构图的支持，还复现并开源了相关异构图神经网络的代码：

△DGL 0.4中基于异构图的RGCN层实现代码

官方表示，在这一新版本的DGL上实现业内比较突出的异构图神经网络，性能上也有更好的表现：

GCMC：DGL的实现相比原作者实现在MovieLens-100K上有5倍加速，在MovieLens-1M上有22倍加速。DGL的内存优化支持在一块GPU上对MovieLens-10M进行训练（原实现需要从CPU动态加载数据），从而将原本需要24小时的训练时间缩短到了1个多小时。

RGCN：使用全新的异构图接口重新实现了RGCN。新的实现大大减少了内存开销。原实现在AM数据集上（边数>5M）由于内存开销太大，只能在CPU上计算，而DGL能使用GPU进行加速并获得了291倍的速度提升。

HAN：提供了灵活的接口可以将一个异构图通过元路径（metapath）转变成同构图。

Metapath2vec：新的元路径采样实现比原作者的实现快两倍。

而且， DGL也发布了针对分子化学的模型库DGL-Chem，以及训练知识图谱嵌入（Knowledge Graph Embedding）专用包DGL-KE。DGL-Chem提供了包括分子性质预测和分子结构生成等预训练模型。

在单GPU上，DGL-KE能在7分钟内使用经典的TransE模型训练出FB15K的图嵌入。相比GraphVite（v0.1.0）在4块GPU上需要14分钟。

DGL-KE首个版本发布了TransE、ComplEx和Distmult模型，支持CPU训练、GPU训练、CPU和GPU混合训练以及单机多进程训练。

DGL团队表示，在未来也会发布更多模型和训练方式。

传送门

DGL异构图教程：

https://docs.dgl.ai/tutorials/hetero/1_basics.html

DGL v0.4 更新笔记：

https://github.com/dmlc/dgl/releases

DGL-KE代码及使用说明：

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

DGL-Chem 模型库:

https://docs.dgl.ai/api/python/model_zoo.html#chemistry

— 完 —

诚挚招聘

量子位正在招募编辑/记者，工作地点在北京中关村。期待有才气、有热情的同学加入我们！相关细节，请在量子位公众号(QbitAI)对话界面，回复“招聘”两个字。

量子位 QbitAI · 头条号签约作者

?'?' ? 追踪AI技术和产品新动态