CogDL:一个全面的图深度学习库

在当今数据驱动的世界中,图数据结构无处不在,从社交网络到生物信息学,再到推荐系统。为了更好地理解和利用这些复杂的关系数据,图深度学习技术应运而生,并迅速成为研究热点。然而,开发和应用图深度学习模型仍然面临诸多挑战,如数据处理、模型设计、训练优化等。为了解决这些问题,清华大学等机构的研究人员开发了CogDL(Cognitive Deep Learning)——一个全面而强大的图深度学习库。

CogDL的主要特点

CogDL作为一个开源的图深度学习工具包,具有以下几个突出特点:

高效性: CogDL利用经过优化的算子来加速图神经网络(GNN)模型的训练,并节省GPU内存。这使得研究人员能够更快地进行实验和迭代。
易用性: CogDL提供了简单易用的API,允许用户通过几行代码就能运行实验,对给定的模型和数据集进行超参数搜索。这大大降低了入门门槛,使得更多人能够参与到图深度学习的研究中来。
可扩展性: CogDL的设计使得将GNN模型应用到新场景变得容易。用户可以基于CogDL的框架快速开发和测试自己的模型。

CogDL logo

CogDL支持的主要任务

CogDL支持图深度学习领域的多项关键任务,包括但不限于:

节点分类
图分类
链接预测
社区检测
网络表示学习
图生成
图自监督学习

对于这些任务,CogDL提供了丰富的数据集和模型实现,使得研究人员可以方便地进行实验对比。

使用CogDL的优势

丰富的模型库: CogDL实现了大量最先进的图深度学习模型,如GCN、GAT、GraphSAGE等,用户可以直接调用这些模型进行实验。
统一的训练接口: CogDL设计了统一的训练循环,引入了DataWrapper和ModelWrapper来简化训练、验证和测试过程。
自动化实验: CogDL的experiment API支持自动化的超参数搜索,帮助用户更轻松地找到最佳模型配置。
高性能计算: CogDL提供了优化的稀疏矩阵乘法操作(GE-SpMM),可以显著加速GNN在GPU上的训练。
灵活的扩展性: 用户可以轻松地将自定义模型集成到CogDL中,利用框架提供的各种功能。

快速开始使用CogDL

安装CogDL非常简单,只需要通过pip执行以下命令:

pip install git+https://github.com/thudm/cogdl.git

安装完成后,您可以使用以下代码快速运行一个实验:

from cogdl import experiment

# 基础用法
experiment(dataset="cora", model="gcn")

# 设置其他超参数
experiment(dataset="cora", model="gcn", hidden_size=32, epochs=200)

# 在多个模型和不同种子上运行
experiment(dataset="cora", model=["gcn", "gat"], seed=[1, 2])