在IPU上启用Graphium，大规模训练分子GNN

Graphium是一个专为在化学任务上进行图表征学习而设计的开源库。

通过可扩展的API，Graphium^[1]提供了先进的GNN架构，使研究人员能够轻松构建和训练自己的大规模GNN。

从发布开始，Graphium就能够充分利用Graphcore（拟未） IPU独特的架构，该架构已被商业和科学研究机构证明可以大大加速GNN的工作负载。

用户可以在Paperspace Notebook上使用ToyMix数据集体验Graphium的功能。有关该notebook的更多信息，请参阅我们的文章^[2]。

https://ipu.dev/8DKBQ0

近年来，伴随着图神经网络（GNN）在社交网络、路线图、药物研发等各种应用中的广泛使用，图表征学习的采用出现了明显的增长。

由于其具备能够更好捕捉分子复杂性质的潜力，用图来表示分子已经越来越受欢迎。在这种方法中，原子被表示为节点，共价键被表示为边，可以用于下游的任务，例如性质预测、化学反应预测和分子生成。因此，GNN已经成为学习这些分子表征的领先架构。

为了在分子图上训练机器学习模型，传统的方法是根据图的局部结构信息提取二进制指纹向量，然后将其输入模型。然而，这些方法有其固有的局限性，如在处理大数据集时的比特碰撞，以及假设分子性质仅取决于局部结构的强烈归纳偏向。

随着深度学习在计算机视觉和自然语言处理等复杂任务中的进展，下一步自然是对分子图采用类似的算法。

通过对最先进的（SOTA）GNN采用这种方法，我们的目标是用能够适应不同任务的学习图嵌入来取代预先定义的指纹。这种学习表征有效地消除了局部结构的归纳偏向。

超大的多层GNN可能会改变我们与分子打交道的方式，因为它们可以学习图连接以外的特征，这些特征不仅嵌入了它们的拓扑结构，而且还嵌入了它们的量子力学和生物性质。

该领域的近期进展重塑了我们对GNN设计的理解，并强调了各个组成部分的重要性，包括位置编码的类型、信息传递层、Transformers中的边缘意识注意力、特征的跨层迁移（节点、边、图、节点对），以及在每一层进行预测的能力。

这是Graphium的真正优势所在：它以灵活的方式整合了所有这些进展，实现了对GNN设计空间的有效探索，并在基于图的工作流程中实现了前所未有的速度、规模和简单性。

Graphium是一个开源的Python库，专注于大规模分子图表征学习。Graphium具有内置的、出色的原子/键特征化工具（featurizer）和模块化的、可配置的API，使任何人（无论他们过去在图学习方面的专业知识如何）都可以轻松地训练分子GNN，以用于分子性质预测、优化和设计等应用。通过为多种类型的硬件优化库，并实施多级多任务，Graphium正在帮助全世界的研究人员释放GNN的真正潜力。

Graphium由研究人员创建，又服务于研究人员，这些研究人员活跃在图机器学习和药物研发的交叉领域。Graphium是唯一一个专门用于在分子图上大规模训练GNN的库，包括以下特点：

 🤑 丰富的特征化：Graphium具有强大而灵活的内置分子特征化工具，用于性质预测。与假设预定义分子特征集不同，Graphium用户可以根据任务的重要性输入他们认为最重要的特征。它提供了38个原子特征（质量、电负性、原子序数、组、周期等）和7个键特征（立体、共轭、长度等）。此外，Graphium还支持构象性质，可以使用高斯核或其他先验知识处理3D位置。

🎨 强大功能：通过配置文件，Graphium可以让那些最佳GNN模型和最新GNN模型的实现变得非常简单，并具有足够的灵活性以满足研究需求。用户可以访问基线和SOTA的2D和3D GNN架构，这包括GCN、GIN、Gated-GCN、PNA、DimeNet、Graphormer以及2022年OGB-LSC获奖模型GPS++^[3]等。

🕹️ 易于使用：通过可扩展的API，研究人员可以轻松构建和训练自己的GNN模型。通过使用配置文件，用户还可以定义模型和训练器的所有参数。

📖 灵活：通过分离GNN层与模型架构，Graphium提供了出色的代码灵活性，可以无缝地替换不同类型的层作为超参数。用户可以使用任何类型的GNN层构建完整的GNN。例如，将基线GCN替换为GPS++只需在YAML文件中更改几行代码。这种增强的灵活性使得Graphium可以通过持续高效地集成新的架构、GNN层或池化层来不断发展和演进。通过分离这些选项的影响，它还支持用户交叉验证参数。

⚡ 快速：基于PyTorch-Lightning库，我们提供了一个简单的训练循环，可以无缝处理数据集、设备类型和指标记录。所有训练参数（学习率、epoch数量、提前停止）和要跟踪的指标都可以在YAML文件中进行配置。其中最佳的模型会自动被保存，并已准备好加载推理任务或迁移学习。Graphium还提供针对GNN的超快速训练。它支持优化的数据处理、数据加载，并且在GPU之外还支持IPU加速器。IPU可以加快稀疏矩阵乘法，实现更快、更具成本效益的GNN训练。

🔥 规模：Graphium有一个具有超过8600万个独特分子的数据集^[4]，每个分子有多达3300个标签。Graphium提供大规模、多任务监督的分子数据集，这些数据集比任何其他公开可用的数据集都要大。通过同时对数百万个分子和数千个任务进行训练，用户可以训练出在化学空间和任务中具有更好通用性的模型。

🦜 表达性： 位置编码是GNN表达性中一个活跃的研究领域，对图Transformers尤其重要。和其他开源库相比，Graphium支持更多的位置编码，包括特征向量、随机游走、通勤距离、热核和图距离。此外，所有这些位置编码在图、节点、边和配对的层面上都得到了支持，这得益于Graphium通过池化/非池化从一个层面转移到另一个层面的新颖能力。所有的位置编码都可以传递给专门设计的编码器，如MLPs、SignNet和高斯核，然后再追加到图的特征中。通过一个库，您可以更好地处理具有表达力的GNN和Transformers的空间。

https://ipu.dev/8DKBQ0

首先，从软件下载中心^[5]下载适合您的操作系统的Poplar SDK，然后按照下面的安装说明进行安装。

mkdir ~/.venv                               # Create the folder for the environment

python3 -m venv ~/.venv/graphium_ipu        # Create the environment

source ~/.venv/graphium_ipu/bin/activate    # Activate the environment

# Installing the poptorch SDK. Make sure to change the path

pip install PATH_TO_SDK/poptorch-3.2.0+109946_bb50ce43ab_ubuntu_20_04-cp38-cp38-linux_x86_64.whl

# Activate popular SDK.

source PATH_TO_SDK/enable

# Install the IPU specific and graphium requirements

PACKAGE_NAME=pytorch pip install -r requirements_ipu.txt

pip install -r lightning.txt

# Install Graphium in dev mode

pip install --no-deps -e .

如欲了解如何训练模型，您可以查看文档或使用Paperspace上现成的模型。

我们相信通过扩展GNN的适用领域，可以将其广泛应用于药物研发的其他任务中，从而产生更广泛的影响。随着我们不断扩展Graphium，我们希望通过前所未有的规模和速度，支持超大规模对接、分子动力学和蛋白质设计等方面的新任务。

如欲获取最新资讯，敬请关注datamol.io^[6]的推特账号。

您可以在这里查看我们的文档^[7]，或在线尝试Graphium^[8]。我们欢迎您在Github存储库^[9]或论坛^[10]留下反馈。您还可以加入Learning on Graphs and Geometry（LoGG）^[11]Slack社区^[12]，在那里用户可以分享他们的成果、提出他们的问题，并在使用Graphium的项目上进行合作。

Graphium库是Valence Discovery、Graphcore以及MILA^[13]的学术人员共同合作的成果。Graphcore团队包括Alex Cunha、Blazej Banaszewski、Zhiyi Li、Sam Maddrell-Mander、Chad Martin、Dominic Masters和Callum Mclean。

本文由Dominique Beaini撰写。Dominique Beaini是蒙特利尔大学信息与运筹学系的兼职教授（Adjunct Professor）、Mila的Associate Industry Member ，以及Valence（最近被Recursion收购）图研究方面的科学家（Staff Scientist）。