支持TensorFlow、PyTorch、Paddle与MindSpore的图神经网络算法库GammaGL

本文将介绍三个方面内容：

• 根据背景介绍GammaGL的两大特点：支持多后端和类PyG的设计原则；

• GammaGL整体架构；

• GammaGL使用方法。

启智社区开源地址：https://git.openi.org.cn/GAMMALab/GammaGL

GitHub开源地址：https://github.com/BUPT-GAMMA/GammaGL

背景介绍

当前，深度学习框架主要由国外的TensorFlow、PyTorch主导，其深度学习生态由欧美主导，优先支持了Nvidia的GPU与Google的TPU，几乎不支持国内的计算设备与操作系统。近年来，国内的AI框架如PaddlePaddle、MindSpore也获得了巨大的发展，以及百度昆仑、华为Ascend等优秀国产芯片国产AI芯片获得了更多的使用。而百花齐放的图深度学习工具往往也是基于现有的深度学习框架，支持的后端也不尽相同。例如，

• 唯一支持多框架的Deep Graph Library (DGL)支持了PyTorch、TensorFlow与MXNet；

• 基于PyTorch开发的PyTorch Geometric (PyG)；

• 受PyG启发的TensorFlow Geometric (TF_Geometric) 以及TensorFlow官方正处于开发阶段的TensorFlow GNN；

• 基于Paddle的Paddle Graph Learning (PGL)；

• 基于MindSpore的MindSpore Graph Learning。

不同深度学习框架有着各自的特点与优势，比如PyTorch凭借着动态图的设计原则在研究领域取得了主导地位，TensorFlow则凭借其效率获得了工业界的青睐。开发者为了运行几个开源算法，需要安装多个框架多个版本的环境、cuda等依赖库版本冲突，同时，国产框架的学习成本高、生态资源少，也一定程度上阻碍了国产AI框架与硬件的发展。

为此，我们开源了支持多深度学习框架的图神经网络算法库，能够支持TensorFlow、PyTorch、PaddlePaddle、MindSpore等作为后端。

同一套代码支持不同深度学习后端

GammaGL能够只使用一套代码在不同深度学习后端运行，这也是与DGL支持多后端的区别所在。

DGL is framework agnostic.

DGL is framework agnostic, meaning if a deep graph model is a component of an end-to-end application, the rest of the logics can be implemented in any major frameworks, such as PyTorch, Apache MXNet or TensorFlow.

DGL利用DLPack在不同框架之间共享张量的结构，能够将整个端到端的GNN模型（通常包含图上的查询与计算等）做到框架无关 (famework agnostic)。然而，一个完整的GNN应用还包含数据加载、GNN模型构建、训练流程、评测等，这些均是框架特定的。例如，使用TensorFlow来替换PyTorch代码时，会涉及到三类变动：

1. 模型类的继承（使用 tensorflow.keras.layer.Layer 替换 torch.nn.Module）

2. 模型内子模块与参数的初始化；

3. 框架特定的操作（如使用 tensorflow.matmul 替换 torch.matmul）

为了避免不同框架下的代码重复编写，GammaGL基于支持多后端的深度学习框架TensorLayerX来实现图神经网络的基本组件与相关算法，将上述三类变动基于TensorLayerX来统一实现，即使用一套代码即可在不同后端运行。

除此之外，用户可能因为框架的某一特性对其有特殊偏好，GammaGL也支持使用框架自身的接口来编写框架特定的操作，即TensorLayerX与某一框架混用的方式。

类PyG的使用方法与用户体验

PyG在Github已获得14.8k stars，无疑是最受欢迎的图神经网络框架之一。同时，PyG是PyTorch-on-the-rocks，即使用Tensor-centric的API，并接近原生PyTorch的设计原则，使得其与PyTorch一样拥有较高的易用性。所以，PyG也积累了大量的用户，顶会源码中也少不了PyG的身影。

因此，GammaGL借鉴了PyG的设计框架，如Tensor-centric的API设计、类torchvision的数据变换与增强等。如果你是一个PyG的用户或非常熟悉PyG，GammaGL将会非常友好，也将会扩展成为TensorFlow Geometric、Paddle Geometric、MindSpore Geometric。

GammaGL整体架构

GammaGL提供了一个多层的框架，为用户提供了低级与高级的接口来构建GNN应用。它由下列组件组成：

• 底层基于深度学习框架来支持，提供了自动微分的计算算子、数据管理与处理流程、训练流程等，并由TensorLayerX使用统一接口进行封装。

• 图数据组件负责图数据的抽象、数据集的管理、图变换等，为上层接口提供存储与查询功能。

• 消息传递组件包含了消息传递机制三大函数的抽象以及底层算子的统一封装，即将常用的scatter_xxx (segment_xxx) 与gather等进行统一并封装。

• 采样组件负责扩大可训练图的规模，通过采样的方式获取子图，进而使用小批量的训练方式。

• 模型层面，GammaGL预先提供了一些可以直接运行的样例模型，还支持用户自定义构建自己的模型。

如何使用GammaGL

快速开始

1. 安装一种你希望使用的后端

      # For tensorflow
      pip install tensorflow-gpu # GPU version
      pip install tensorflow # CPU version
   
      # For torch, version 1.10
      # https://pytorch.org/get-started/locally/
      pip3 install torch==1.10.2
   
      # For paddle, any latest stable version
      # https://www.paddlepaddle.org.cn/
      python -m pip install paddlepaddle-gpu
   
      # For mindspore, GammaGL only supports version1.6.1, GPU-CUDA 11.1 and python 3.7.5
      # https://www.mindspore.cn/install
      pip install https://ms-release.obs.cn-north-4.myhuaweicloud.com/1.6.1/MindSpore/gpu/x86_64/cuda-11.1/mindspore_gpu-1.6.1-cp37-cp37m-linux_x86_64.whl --trusted-host ms-release.obs.cn-north-4.myhuaweicloud.com -i https://pypi.tuna.tsinghua.edu.cn/simple
   

2. 安装TensorLayerX

   pip install git+https://github.com/tensorlayer/tensorlayerx.git 
   # 也可从OpenI启智社区安装，更适合大陆用户
   # pip install git+https://git.openi.org.cn/OpenI/TensorLayerX.git     
   

3. 安装GammaGL

   git clone https://github.com/BUPT-GAMMA/GammaGL.git
   # 也可从OpenI启智社区clone，更适合大陆用户
   # git clone https://git.openi.org.cn/GAMMALab/GammaGL.git
   python setup.py install
   

我们在GitHub主页提供了快速构建一个GNN模型与训练流程的样例。

除此之外，GammaGL还提供了二十余个算法供大家使用，包含同质图上一些经典的模型、对比学习模型、异质图神经网络模型等。

同样地，我们仅需使用命令行即可运行上述模型。

以GCN为例，仅需指定环境变量即可运行相应的后端：

# cd ./examples/gcn
# set parameters if necessary
CUDA_VISIBLE_DEVICES=0 TL_BACKEND=paddle python gcn_trainer.py

还能通过脚本中的os模块来指定后端

import os
os.environ['TL_BACKEND'] = 'paddle' # tensorflow, torch, mindspore

写在最后

GammaGL在GitHub与启智社区同时开源并同步更新，接入TensorLayerX开源生态，积极为国产AI生态贡献自己的绵薄力量。同时，GammaGL将持续不断投入，加快新功能的开发与版本的迭代：聚焦于更多算法模型的集成、常用操作的开发、算子的优化与迭代等。

GammaGL欢迎大家使用以及提出宝贵的意见，也欢迎各位开源贡献者参与、共建社区生态。任何问题欢迎联系: [email protected]

传送门：Tensorlayer团队的最新作品，跨平台开源框架TensorLayerX来了！