当前位置: 首页 > 科技观察

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

时间:2023-03-16 23:34:06 科技观察

开源图神经网络框架DGL升级:GCMC训练时间从1天缩短至1小时又一AI框架迎来升级。此次纽约大学与亚马逊联合推出图神经网络框架DGL。不仅全面推出了对异构图的支持,相关异构图神经网络的代码也进行了复现和开源,并在GCMC、GCMC等业界知名模型的实现中取得了较好的效果。RCGN。还发布了训练知识图谱嵌入(KnowledgeGraphEmbedding)的专用包DGL-KE,在很多经典图嵌入模型上实现了性能提升。目前,这个在2018年NeurIPS大会上发布的框架已经获得了3000+的star,已经是图神经网络和图机器学习领域比较受关注的工具。异构图:神经网络的新方向异构图是一个与同构图相对应的新概念。传统的HomogeneousGraph数据中只有一种节点和边,因此在构建图神经网络时,所有节点共享相同的模型参数和相同维度的特征空间。然而,异构图中可以有不止一种节点和边,因此允许不同类型的节点具有不同维度的特征或属性。这种特性使得异构图的应用非常广泛。如果我们用图来描述我们与周围事物的关系,我们会发现生成的图自然是异构的。比如这个例子:今天看了电影《流浪地球》,作为观众的“我”和电影《流浪地球》建立了“看”的关系。异构图可用于描述此类交互的集合。这张图分为“观众”和“电影”两种类型的节点,以及一条“观看”的边。作为观众,“我”必须具有与电影不同的属性,需要用不同的模型或不同的特征维度来表达。所以,这张图自然是异类。而且,与传统方法相比,基于异构图训练的神经网络在某些场景下也能取得更好的效果和性能。目前已应用于知识图谱、推荐系统、恶意账号识别等领域和任务。两个代表性模型是用于节点分类和链接预测的RGCN和用于产品推荐的GCMC。但是如何设计一个“又快又好”的深度神经网络呢?这是纽约大学与亚马逊联合推出图神经网络框架DGL的起点。新版DGL:异构图神经网络工具本次更新发布了DGL0.4版本,让整个框架更加实用。不仅全面推出了对异构图的支持,还转载并开源了相关异构图神经网络的代码:△DGL0.4基于异构图的RGCN层的实现代码官方表示是在这个新版本的DGL业界优秀的异构图神经网络也有更好的性能:GCMC:与原作者对DGL的实现相比,DGL的实现在MovieLens-100K上快了5倍,在MovieLens上快了22倍-1M。DGL的内存优化支持在GPU上训练MovieLens-10M(原先的实现需要从CPU动态加载数据),从而将原本需要24小时的训练时间缩短为1小时以上。RGCN:RGCN使用全新的异构图形接口重新实现。新的实现大大减少了内存开销。最初的实现是在AM数据集上执行的(边数>5M)。由于内存开销大,只能在CPU上计算,而DGL可以使用GPU加速,获得291倍的速度提升。HAN:提供灵活的接口,通过元路径将异构图转换为同构图。Metapath2vec:元路径采样的新实现比原作者的实现快一倍。此外,DGL还发布了分子化学模型库DGL-Chem,以及训练知识图谱嵌入(KnowledgeGraphEmbedding)的专用包DGL-KE。DGL-Chem提供预训练模型,包括分子特性预测和分子结构生成。在单GPU上,DGL-KE可以使用经典的TransE模型在7分钟内训练出FB15K的graphembedding。与在4个GPU上需要14分钟的GraphVite(v0.1.0)相比。DGL-KE第一版发布了TransE、ComplEx和Distmult模型,支持CPU训练、GPU训练、CPU和GPU混合训练、单机多进程训练。DGL团队表示,未来会发布更多的模型和训练方法。