GCPNs模型使用TorchDrug从头生成分子

GCPNs模型使用TorchDrug从头生成分子

新分子的生成已经成为寻找新疾病治疗方法的最受欢迎的方法之一，不出所料，我们看到许多分子生成都是使用机器学习完成的。当然，对于分子，有一些硬性规定:药物分子必须具有一定的特性才能在现实世界中存在。

这就是强化学习的作用所在。

在这里，重点强调图卷积策略网络(GCPNs)是很重要的。gcpn于2018年首次提出，是图神经网络和强化学习应用于药物发现问题的基石。利用RL，我们可以优化分子图生成过程，添加满足我们要求的特定子结构。

然后TorchDrug的好示例出现了。他们创建了一个一站式框架，用于实现大量基于图的模型，以解决各种各样的问题。我们今天关注的是gcpn的分子生成。

在本报告中，我们将带您浏览分子图，图卷积策略网络是如何工作的，最后使用TorchDrug训练我们自己的模型，并使用Weights & Biases帮助记录。让我们开始

关于分子图的快速构建

首先，如果你需要快速复习一下分子生物学:分子实际上可以用异质图表示。在这里，原子相当于节点，键相当于边。图可以由$d$类型的节点和$b$类型的边组成。一个有$n$个节点的图可以用一组顶点 F i n { 0 , 1 } n × d F in \{0,1\}^{n \times d} Fin{0,1}n×d和邻接张量 E ∈ { 0 , 1 } n × n × b E \in \{0,1\}^{n \times n \times b} E∈{0,1}n×n×b表示。

Source: Improving graphs of cycles approach to structural similarity of molecules. PLoS ONE: 14(12)

图卷积策略网络

图形生成

图形生成过程被建模为一系列图形编辑：选择两个原子并在它们之间构建一条边。这看起来很简单，但我们关心的问题是我们选择哪一对原子?既然我们有多种类型的边，我们应该用哪一种？

这就是我们希望模型学习的东西!因此，GCPN模型在每一步使用图卷积来依次计算四个概率

1. 从构造的图中挑选每个原子的概率。
2. 从迄今为止构建的任一图中选择另一个原子的概率, 或者使用上一步中选择的原子在我们的原子类型词汇表中的 d 个原子中选择另一个原子的概率。
3. 使用每种边类型连接前两步中选择的原子的概率。
4. 停止生成过程的概率。
   
   Source: Graph Convolutional Policy Network for Goal-Directed Molecular Graph Generation. NeurIPS 2018: 6412-6422

奖励

如果生成的分子不违反任何价性规则，则每一步都有中间奖励;根据优化模型的属性，当生成结束时，还会有最终奖励。在这份报告中，我们将看到的属性是被惩罚的LogP (pLogP)，这是一种度量

• 将药物运送到体内将有多容易/困难。
• 合成实验室中的分子会多么容易/困难。

在本报告中，我们尝试最大化$PLOGP$。

预训练

该模型首先在大量的类药物分子上进行预训练，以便学习分子生成的一般规则。向该模型提供数据集中图的随机子图，然后期望该模型重构该图。对于本报告，ZINC250K数据集的一个较小子集(包含10K分子)用于此任务。最后给出了使用W&B进行测TorchDrug训练引擎的实现。

下一个代码片段加载数据集和GCPN策略，并使用新的训练引擎对其进行50个epochs预训练，该引擎允许在W&B上记录度量，并在每个阶段将模型保存为artifact，以便为下一步加载。

它还从生成的50个分子中创建了一个包含3个分子的表，用于验证以及相应plogp的分布和平均值。

from torchdrug import datasets,models, tasks
from torch import optim


dataset = datasets.ZINC250k("~/molecule-datasets/", kekulize=True,
                            node_feature="symbol")
model = models.RGCN(input_dim=dataset.node_feature_dim,
                    num_relation=dataset.num_bond_type,
                    hidden_dims=[256, 256, 256, 256], batch_norm=False)
task = tasks.GCPNGeneration(model, dataset.atom_types, max_edge_unroll=12,
                            max_node=38, criterion="nll")
optimizer = optim.Adam(task.parameters(), lr=1e-3)
solver = Engine(task, dataset, None, None, optimizer, gpus=(0,), batch_size=128, log_interval=10)
solver.train(num_epoch=50)

在训练过程的开始，我们看到生成的分子几乎完全由碳原子组成，但当我们接近训练过程的结束时，我们看到了属于类药物分子的亚结构的添加，如氮原子、氧原子和硫原子的存在。

优化

一旦生成器模型被预训练为通用分子，下一个任务是优化它，以生成具有高pLogP的分子。为此，利用近端策略优化**(PPO)**优化图卷积策略的权值，使奖励最大化。

为了做到这一点，我们首先通过上一步加载最新训练的模型artifact。

api = wandb.Api()


artifact = api.artifact("GCPN-pLogP/gcpn:latest")
artifact_dir = artifact.checkout()


model = models.RGCN(input_dim=18,
                    num_relation=3,
                    hidden_dims=[256, 256, 256, 256], batch_norm=False)   
task = tasks.GCPNGeneration(model, [6, 7, 8, 9, 15, 16, 17, 35, 53],
                            max_edge_unroll=12, 
                            max_node=38, 
                            task="plogp", 
                            criterion="ppo",
                            reward_temperature=1,
                            agent_update_interval=3, 
                            gamma=0.9)

optimizer = optim.Adam(task.parameters(), lr = 1e-5)
solver = Engine(task, dataset1, None, None, optimizer,
                     gpus=(0,), batch_size=512, log_interval=10)

solver.model.load_state_dict(torch.load("./artifacts/gcpn:latest/model_gcpn.tar"))

现在我们已经有了一个装载了预训练的GCPN的模型，我们可以继续进行RL优化

solver.train(num_epoch=20, train_type='rl', wandb_logger=True)

在每个epoch结束，我们生成50个分子并记录其plogps以验证该模型的性能。随着epoch增加，plogps的直方图被认为朝向更积极的区域，这表明该模型学习了应添加哪些子结构以最大程度地提高属性。

结论

本报告中展示的可视化使Weights & Biases在药物发现工作流中成为一个非常棒的工具，并有助于对模型训练阶段提供巨大的洞察力。此外，artifacts使用提供了一种简单的方法来记录数据集以及训练好的模型，以便在不同的下游任务中重用，还有助于确保可以一致地重新创建实验。

参考

https://wandb.ai/manan-goel/GCPN/reports/De-Novo-Molecule-Generation-with-GCPNs-using-TorchDrug–VmlldzoxNDgzMzQz