基于图深度学习的开源工具包——CogDL

一行代码命令可以做什么？

“一行命令可以实现‘一条龙’运行实验。”

访问 github.com/THUDM/cogdl 一键体验！

近年来，结构化数据的表示学习备受业界关注与热捧，图神经网络成为处理相关工作的有力工具，基于随机游走、矩阵分解的方法在搜索推荐、分子和药物生成等领域有着十分重要的应用。

但是，由于许多项目的代码并未开源或者开源代码的风格多种多样，研究者和使用者在使用这些方法的过程中会遇到各种各样的问题，比如实验复现以及如何在自己的数据集上运行模型等。

具体而言（以学术研究为例），研究者和使用者在从事研究过程中需要将自己提出的模型与其他模型进行对比，来验证所提出模型的有效性。但在对比过程中，他们需要选定若干个下游任务，在每个任务下，通过公平的评估方式来对比不同模型的性能，由于不同模型在提出时可能会使用不完全一致的下游任务或者评估方式，他们需要花费大量精力修改基线模型来进行适配。

那么，如何快速、便捷地复现出基线模型（baseline）的结果，并将这些模型应用到自定义的数据集上？

为此，清华大学知识工程实验室（KEG）联合北京智源人工智能研究院（BAAI）开发了一种基于图深度学习的开源工具包——CogDL（底层架构为 PyTorch，编程语言为 Python）。

据 CogDL 开发者介绍，该工具包通过整合多种不同的下游任务，同时搭配合适的评估方式，使得研究者和使用者可以方便、快速地运行出各种基线模型的结果，进而将更多精力投入研发新模型的工作之中。

“对图领域中每种任务，我们提供了一套完整的“数据处理-模型搭建-模型训练-模型评估”的方案，易于研发人员做相关的实验。比如对于图上半监督节点分类任务，我们整合了常用的数据集 Cora、Citeseer、Pubmed，提供了经典的/前沿的各种模型（包括GCN、GAT、GCNII 等），提供了相应的训练脚本，并且整理出了一个相应的排行榜作为参考。” CogDL 开发者说。

CogDL 最特别的一点在于它以任务（task）为导向来集成所有算法，将每一个算法分配在一个或多个任务下，从而构建了 “数据处理-模型搭建-模型训练和验证” 一条龙的实现。

此外，CogDL 也支持研究者和使用者自定义模型和数据集，并嵌入在 CogDL 的整体框架下，从而帮助他们提高开发效率，同时也包含了当前许多数据集上 SOTA 算法的实现，并且仍然在不断更新。 CogDL：面向任务，扩展算法 图表示学习算法可以分为两类：一类是基于图神经网络的算法，另一类是基于 Skip-gram 或矩阵分解的算法。前者包括 GCN、GAT、GraphSAGE 和 DiffPool 等，以及适用于异构图的 RGCN、GATNE 等；后者则包括 Deepwalk、Node2Vec、HOPE 和 NetMF 等，以及用于图分类的 DGK、graph2vec 等算法。

大体上，CogDL 将已有图表示学习算法划分为以下 6 项任务：

有监督节点分类任务（node classification）：包括 GCN、GAT、GraphSAGE、MixHop 和 GRAND 等；

无监督节点分类任务（unsupervised node classification）：包括 DGI、GraphSAGE（无监督实现），以及 Deepwalk、Node2vec、ProNE 等；

有监督图分类任务（graph classification）：包括 GIN、DiffPool、SortPool 等；

无监督图分类任务（unsupervised graph classification）：包括 InfoGraph、DGK、Graph2Vec 等；

链接预测任务（link prediction）：包括 RGCN、CompGCN、GATNE 等；

异构节点分类（multiplex node classification）：包括 GTN、HAN、Metapath2vec 等。

CogDL 还包括图上的预训练模型 GCC，GCC 主要利用图的结构信息来预训练图神经网络，从而使得该网络可以迁移到其他数据集上，来取得较好的节点分类和图分类的效果。 能用来做什么？ 那么，研究者和使用者可以利用 CogDL 做些什么？主要有三点：跟进 SOTA、复现实验，以及自定义模型和数据。

跟进 SOTA。CogDL 跟进最新发布的算法，包含不同任务下 SOTA 的实现，同时建立了不同任务下所有模型的 leaderboard（排行榜），研究人员和开发人员可以通过 leaderboard 比较不同算法的效果。

复现实验。论文模型的可复现性是非常重要的。CogDL 通过实现不同论文的模型也是对模型可复现性的一个检验。

自定义模型和数据。“数据-模型-训练”三部分在 CogDL 中是独立的，研究者和使用者可以自定义其中任何一部分，并复用其他部分，从而提高开发效率。 怎么用？ 有两种方法：命令行直接运行，通过 API 调用 命令行直接运行。通过命令行可以直接指定 "task"、"model"、"dataset" 以及对应的超参数，并且支持同时指定多个模型和多个数据集，更方便。代码如下：

支持自定义数据集和模型，并且提供了自动调参的例子。

据 CogDL 研究者介绍，从 2019 年启动以来，该项目经过了多次密集的研发，到目前已经整合了图领域的重要任务及其相关的经典和前沿的模型。

CogDL 开发者表示，“起初，我们只是整合了实验室内部从事相关领域研究的同学所使用的相关代码，为了方便实验室其余的同学运行相关的实验。后来，我们决定将代码开源出来，让更多的研究和开发人员能够借助我们的代码从事相关工作。”

如今，CogDL 已经上线了 CogDL v0.1.1。

“我们增加了一些最前沿的图神经网络模型，包括图自监督模型、图预训练模型。用户可以利用图上预训练好的模型做不同的下游任务。在下游任务方面，我们增加了知识图谱的链接预测任务，方便从事知识图谱领域的用户进行相关实验。” CogDL 开发者说。

具体而言，在监督节点分类任务上，增加了 GRAND 和 DisenGCN，GRAND 在 Cora、PubMed、Citeseer 上取得了 SOTA 或者接近 SOTA 的效果；无监督节点分类任务上，增加了 DGI、MVGRL、GraphSAGE（无监督实现）；在链接预测任务上，增加了知识图谱上的链接预测 GNN 模型 RGCN 和 CompGCN；以及增加了图上的预训练模型 GCC。

另外，CogDL 也提供了使用 optuna 进行超参数搜索的 example，通过指定 "模型，数据集，参数" 即可自动实现超参数的搜索。

接下来，CogDL 研究者希望在现有基础上，继续补充其余的图领域的下游任务，添加更多的数据集，更新每个任务的排行榜；同时增加最前沿的关于预训练的图神经网络模型，支持用户直接使用预训练好的模型来进行相关应用；以及完善 CogDL 的教程和文档，让刚接触图领域的初学者能够快速上手。

谈及最终的愿景，CogDL 研究者表示，一方面希望从事图数据领域的研究人员可以方便地利用 CogDL 来做实验，另一方面希望面向实际应用的开发人员可以利用 CogDL 来快速地搭建相关下游应用，最终打造一个完善的图领域的研发生态。
编辑：hfy