在IPU上使用Graphium进行多任务分子建模

Graphium是一个高性能深度学习库，专门用于各种化学任务的图表征学习。Graphium集成了最先进的图神经网络（GNN）架构和用户友好的API，可以轻松构建和训练自定义的GNN模型。

Graphium的独特特点是其强大的特征提取能力，能够从分子结构中提取全面的特征，适用于各种应用，包括性质预测、虚拟筛选和药物研发。

Graphium适用于广泛的化学任务，如性质预测、毒性评估、分子生成和图回归。无论是处理简单的分子还是复杂的化学图形，Graphium都能提供高效的工具和模型。

更多关于Graphium的信息，请查看我们的文章^[1]。

在这篇文章中，我们将介绍如何使用Paperspace Gradient Notebook展示Graphium在ToyMix数据集上的能力，该数据集是QM9、Tox21和ZINC 12k数据集的组合。我们将深入探讨多任务场景，并展示Graphium如何训练模型通过化学感知同时预测多个性质。

https://ipu.dev/8DKBQ0

总结

领域：分子
任务：多任务
模型：GCN/GIN/GINE
数据集：QM9、Zinc、Tox21
工作流：训练、验证、推理
IPU数量：推荐：4倍（最小：1倍，最大：16倍）
执行时间：20分钟

学习成果

在该Notebook中，您将学习如何：

在IPU上使用Graphium资料库为各种GNN模型运行多任务训练
将结果与单一任务性能进行比较
学习如何修改Graphium的配置文件

其他资源链接

对于该Notebook，如果您已经熟悉GNN和PyTorch Geometric，将会快速上手。

如果您需要快速入门，我们建议您参考在Graphcore（拟未） IPU上使用PyTorch Geometric的教程^[2]，该教程涵盖了GNN和PyTorch Geometric的基础知识，并介绍了如何在IPU上运行模型。

依赖项

我们为这个Notebook安装了以下依赖项（第一次可能需要5-10分钟）：

在Paperspace上运行

Paperspace环境可以让您在没有设置的情况下运行该Notebook。如果问题仍然存在，或者您想向我们提供关于Notebook内容的反馈，请通过Slack^[3]上的Graphcore开发者社区联系我们或者在GitHub^[4]上提问。

要求：

Python packages installed with pip install -r ./requirements.txt

数据集：ToyMix基准

ToyMix数据集是化学机器学习领域中三个被广泛认可的数据集的组合：

QM9
Tox21
ZINC12K

每个数据集都贡献了不同的方面，确保了一个全面的基准测试工具。

QM9在3D GNN领域广为人知，它提供了19个与分子能量最小化3D构象相关的图层面的量子属性。鉴于这些分子最多只有9个重原子，它更简单，并且能够实现快速的模型迭代，就像提出的较大的量子数据集，但规模较小。

Tox21数据集在药物研发的机器学习研究人员中颇受关注。它是一个包含12个标签的多标签分类任务，标签缺失和负类别的不平衡度较高。这反映了较大数据集的特点，特别是在稀疏性和不平衡性方面，使其对于学习如何处理现实世界中的倾斜数据集非常有价值。

最后，ZINC12K数据集在GNN表达性的研究中得到认可，这是大规模数据性能的一个基本特征。ZINK12K被包含在ToyMix中，预计在该任务的成功表现将与更大数据集的成功具有关联。

多任务学习

在多任务学习中，相关的任务在训练期间相互分享信息，从而形成一个共享的表征。这通常可以提高模型在单个任务上的表现，因为从一个任务中学习的信息可以帮助其他任务。在这里，ToyMix数据集提供了一个沙盒，为分子任务开发这些模型。

下载数据集和拆分

ToyMix数据集需要下载三个数据集。这里我们将创建预期的数据目录，并下载原始数据集CSV文件，以及指定训练、验证和测试分割的分割文件。

上述数据集的表征方法表明，QM9和ZINC是密集填充的数据集，具有回归任务和为每个分子提供的每个标签。我们还看到，Tox21数据集是由二进制标签稀疏地填充的。这些数据集的相对大小很重要。QM9大约比其他数据集大约10倍，因此我们预计它会主导训练。

值得注意的是这些数据集的差异，其中分子在尺寸和组成方面来自相当不同的领域。此外，从湿式实验室获得的结果相比，使用DFT计算的分子的量子特性与是非常不同的特性，特别是在复杂系统的内容方面。

为了了解数据集的重叠情况，下面显示了一个统一的模态逼近和投影（UMAP）。在这里我们可以看到分子相似性的二维投影，紫色的QM9结果与Tox21和ZINC结果基本分开。

现在的问题是：大型的量子特性数据集可以在多大程度上用于提高小型数据集的预测能力？

这种用较大的数据集来提高较小的数据集的预测能力的方法，在大型语言模型的世界里产生了令人印象深刻的结果，但在分子方面是否也能做到这一点还不得而知。

Graphium模型的配置

Graphium模型由以下部分组成：

处理分子的输入特征
一组用于位置和结构特征的编码器
进行特征化处理以产生PyG图数据集
由MPNN和转化器层组成的GNN层堆栈--模型的骨干
输出节点、图和边缘嵌入。这些都是由MLPs处理，以产生多任务预测

如下图所示。

可用的配置显示在下面的单元格中，您可以通过取消对所需行的注释在它们之间轻松切换。第一组显示了三种GNN架构在ToyMix数据集上的完整多任务训练。

最后两个配置显示了仅在QM9和Zinc数据集上训练的相同的GCN模型。

更新配置

Graphium库被设计为可以与一个配置文件一起使用。本教程包含三个示例（GCN、GIN和GINE模型），用于多任务ToyMix数据集的训练。

配置文件被分成几组，主要选项包括加速器、数据模块、结构、预测器、指标和训练器。建议您浏览这些文件，了解模型是如何被建构的。

对于某些参数，在notebook中或通过命令行参数来设置它们更加合理。你也可以在notebook中手动设置这些值，如下所示。

(--predictor.torch_scheduler_kwargs.max_num_epochs=200将设置epochs的数量）

在下面的单元格中，我们可以设置复制的数量来加快训练速度。默认值是4，但是如果您使用的是POD16，我们可以增加复制数，对加快训练速度有显著的提升。

我们还可以设置epochs的数量和学习率。对于epochs的数量，YAML配置文件已经被加载到一个dict中，所以我们不能依靠YAML配置文件中的锚和参考系统来保持参数的同步，所以这里要手动完成。

让我们更改数据集的路径以从存储位置获取数据。

关键的超参数：

此处显示的是简单的例子——对于测试，您可以更改此处的值/添加其他参数。当建立一个完整的训练脚本时，我们建议使用配置文件或命令行来设置值。

首先构建预测器——这是模型，查看打印的摘要以了解层类型和参数数量。

然后建立训练器——它控制实际的训练循环。

训练模型

现在我们可以开始训练模型了。训练是由PyTorch Lightning处理，这使得训练循环的设置变得简单。

这将在计算批次中节点和边的最大数量后设置训练循环。我们在IPU上使用固定大小的批次，因此在编译模型之前需要添加一些填充以考虑可变的分子大小。

训练过程显示在进度条上。

小练习： 我们目前只运行了10个epochs的训练来说明训练模型，这样您就可以看到验证和测试的输出格式。尝试增加epochs的数量以获得更好的结果。使用大约30个epochs应该可以得到训练了300个epochs的模型的95%的性能。

验证和测试数据集的拆分

最后，我们可以查看验证和测试数据集的拆分，其中训练的模型从检查点加载并针对保留的数据集进行评估。

结论和下一步计划

我们演示了如何在IPU上使用Graphium进行多任务分子建模。我们使用了ToyMix数据集，该数据集结合了QM9、Tox21和ZINC 12k数据集，并展示了Graphium如何训练模型通过化学感知同时预测多个性质。

从验证和测试结果来看，我们可以看到较大的QM9数据集对改善较小数据集的单一任务性能的影响。准确率和MAE分数比我们在小数据集上训练的预期要好得多，这证实了在非常短的训练运作中，多任务设置也能带为训练带来效益。为了进一步验证这一点，我们可以查看在多任务数据集和单一数据集上训练了全部epoch的同一模型，并比较最终的验证结果。

您可以尝试以下操作：

将训练的epoch数量增加到30左右。这应该会产生相当不错的结果。如果您训练了300个epoch，您的结果应该与基准结果相匹配。
更改配置文件以使用不同的模型。GIN的准确率将比GCN更高，而GINE的准确率将比GIN更高。
调整多任务配置，只使用单个数据集。您可以以此来比较多任务性能与单一任务性能的结果。请记住从数据集、损失和指标中删除相关部分。我们提供了示例YAML配置文件，展示了如何在QM9（config_small_gcn_qm9.yaml）和ZINC 12k（config_small_gcn_zinc.yaml）数据集上使用GCN进行此操作。
在molfeat notebook pytorch_geometric_molfeat.ipynb中探索更多分子特征提取方法。

[1]https://www.graphcore.ai/posts/graphium-an-ipu-ready-python-library-for-training-molecular-gnns-at-scale

[2]https://github.com/graphcore/Gradient-Pytorch-Geometric/tree/main/learning-pytorch-geometric-on-ipus

[3]https://www.graphcore.ai/join-community

[4]https://github.com/graphcore/examples

本篇博客作者：

Sam Maddrell-Mander