跳转到主要内容

分子深度生成模型的实现。

项目描述

MoLeR:分子生成模型

CI license pypi python code style

此存储库包含MoLeR模型的训练和推理代码,该模型在学习使用结构基序扩展分子骨架中介绍。我们还包含了对CGVAE的实现,但不包括与高级模型界面的集成。

快速入门

molecule_generation可以通过pip安装,但它还依赖于rdkit和(如果想要使用GPU)设置CUDA库。两者都可以通过conda获得

conda env create -f environment.yml
conda activate moler-env

我们的软件包已在python>=3.7tensorflow>=2.1.0rdkit>=2020.09.1上进行测试;请参阅environment*.yml文件以了解CI中测试的确切配置。

要安装最新版本的molecule_generation,请运行

pip install molecule-generation

或者,在根目录中运行pip install -e .可以安装代码的最新状态,包括合并到main但尚未发布的更改。

使用默认超参数训练的MoLeR检查点已提供,在此。此文件需要保存到新文件夹 MODEL_DIR(例如,/tmp/MoLeR_checkpoint)并重命名为具有 .pkl 扩展名(例如,到 GNN_Edge_MLP_MoLeR__2022-02-24_07-16-23_best.pkl)。然后您可以运行以下命令来采样10个分子:

molecule_generation sample MODEL_DIR 10

下面是训练您自己的模型和运行更高级推断的说明。

故障排除

问题:在我的系统上安装 tensorflow 不起作用,或者虽然起作用但 GPU 未能使用。

答案:请参阅 tensorflow 网站 的指南。特别是,在 tensorflow 的最新版本中,可能会遇到“找不到 libdevice”的错误;在这种情况下,请按照 此页面底部 的说明操作。

问题:我特定的依赖版本组合不起作用。

答案:请提交问题,并在同时默认使用 environment-py*.yml 中的固定配置之一。

问题:我在中国,因此 figshare 检查点链接对我不起作用。

答案:您可以尝试使用此链接 此链接

工作流程

使用 MoLeR 的工作可以大致分为四个阶段

  • 数据预处理,其中将 plain text 列表中的 SMILES 字符串转换为包含分子图和生成跟踪的 *.pkl 文件;
  • 训练,其中 MoLeR 在预处理数据上进行训练,直到收敛;
  • 推断,其中加载模型并执行批量编码、解码或采样;(可选的)
  • 微调,其中在新的数据上对先前训练的模型进行微调。

此外,您还可以可视化模型的解码跟踪和内部动作概率,这对于调试可能很有用。

数据预处理

要运行预处理,您的数据必须遵循简单的 GuacaMol 格式(文件 train.smilesvalid.smilestest.smiles,每个文件包含 SMILES 字符串,每行一个)。然后,您可以通过运行以下命令来预处理数据:

molecule_generation preprocess INPUT_DIR OUTPUT_DIR TRACE_DIR

其中 INPUT_DIR 是包含三个 *.smiles 文件的目录,OUTPUT_DIR 用于中间结果,TRACE_DIR 用于包含生成跟踪的最终预处理文件。此外,preprocess 命令接受命令行参数以覆盖各种预处理超参数(特别是基团词汇表的大小)。此步骤大致相当于将我们的论文中的算法 2 应用到输入数据中的每个分子。

运行上述命令后,您应看到类似的输出:

2022-03-10 11:22:15,927 preprocess.py:239 INFO 1273104 train datapoints, 79568 validation datapoints, 238706 test datapoints loaded, beginning featurization.
2022-03-10 11:22:15,927 preprocess.py:245 INFO Featurising data...
2022-03-10 11:22:15,927 molecule_dataset_utils.py:261 INFO Turning smiles into mol
2022-03-10 11:22:15,927 molecule_dataset_utils.py:79 INFO Initialising feature extractors and motif vocabulary.
2022-03-10 11:44:17,864 motif_utils.py:158 INFO Motifs in total: 99751
2022-03-10 11:44:25,755 motif_utils.py:182 INFO Removing motifs with less than 3 atoms
2022-03-10 11:44:25,755 motif_utils.py:183 INFO Motifs remaining: 99653
2022-03-10 11:44:25,764 motif_utils.py:190 INFO Truncating the list of motifs to 128 most common
2022-03-10 11:44:25,764 motif_utils.py:192 INFO Motifs remaining: 128
2022-03-10 11:44:25,764 motif_utils.py:199 INFO Finished creating the motif vocabulary
2022-03-10 11:44:25,764 motif_utils.py:200 INFO | Number of motifs: 128
2022-03-10 11:44:25,764 motif_utils.py:203 INFO | Min frequency: 3602
2022-03-10 11:44:25,764 motif_utils.py:204 INFO | Max frequency: 1338327
2022-03-10 11:44:25,764 motif_utils.py:205 INFO | Min num atoms: 3
2022-03-10 11:44:25,764 motif_utils.py:206 INFO | Max num atoms: 10
2022-03-10 11:44:25,862 preprocess.py:255 INFO Completed initializing feature extractors; featurising and saving data now.
 Wrote 1273104 datapoints to /guacamol/output/train.jsonl.gz.
 Wrote 79568 datapoints to /guacamol/output/valid.jsonl.gz.
 Wrote 238706 datapoints to /guacamol/output/test.jsonl.gz.
 Wrote metadata to /guacamol/output/metadata.pkl.gz.
(...proceeds to compute generation traces...)

预处理图保存到 OUTPUT_DIR 后,它们将被转换为具体的生成跟踪,这是预处理中最计算密集的部分。在此部分中,预处理代码可能会打印错误,指出无法解析的分子或失败其他断言;MoLeR 的预处理对这类情况具有鲁棒性,并将简单地跳过任何有问题的样本。

训练

TRACE_DIR 下存储了一些预处理的日期后,MoLeR 可以通过运行以下命令进行训练:

molecule_generation train MoLeR TRACE_DIR

train 命令接受许多命令行参数以覆盖训练和架构超参数,其中大多数通过传递 --model-params-override 来访问。例如,以下命令使用 GGNN 风格的消息传递(而不是默认的 GNN_Edge_MLP)并在编码器和解码器 GNN 中使用更少的层来训练 MoLeR 模型

molecule_generation train MoLeR TRACE_DIR \
    --model GGNN \
    --model-params-override '{"gnn_num_layers": 6, "decoder_gnn_num_layers": 6}'

由于 tf2-gnn 非常灵活,MoLeR 支持广泛的架构配置。

运行 molecule_generation train 后,您应看到类似的输出:

(...tensorflow messages, hyperparameter dump...)
Initial valid metric:
Avg weighted sum. of graph losses:  122.1728
Avg weighted sum. of prop losses:   0.4712
Avg node class. loss:                 35.9361
Avg first node class. loss:           27.4681
Avg edge selection loss:              1.7522
Avg edge type loss:                   3.8963
Avg attachment point selection loss:  1.1227
Avg KL divergence:                    7335960.5000
Property results: sa_score: MAE 11.23, MSE 1416.26 (norm MAE: 13.89) | clogp: MAE 10.87, MSE 4620.69 (norm MAE: 5.98) | mol_weight: MAE 407.42, MSE 185524.38 (norm MAE: 3.70).
   (Stored model metadata and weights to trained_model/GNN_Edge_MLP_MoLeR__2022-03-01_18-15-14_best.pkl).
(...training proceeds...)

默认情况下,训练会在连续 3 个 mini-epochs 中没有改进验证损失时停止,其中 mini-epoch 定义为 5000 个训练步骤;这可以通过 --patience 标志和 num_train_steps_between_valid 模型参数分别进行控制。

推断

模型训练并保存到MODEL_DIR后,我们提供了两种加载方式:通过命令行界面(CLI)或直接从Python加载。目前,基于CLI的加载方式并没有暴露所有有用的功能,主要用于简单的测试。

要使用CLI从模型中采样分子,只需运行

molecule_generation sample MODEL_DIR NUM_SAMPLES

同样,将存储在SMILES_PATH下的SMILES列表编码为潜在向量,并将它们存储在OUTPUT_PATH

molecule_generation encode MODEL_DIR SMILES_PATH OUTPUT_PATH

在所有情况下,MODEL_DIR表示包含模型检查点的目录,而不是检查点的路径本身。模型加载器将只查看MODEL_DIR下的*.pkl文件,并期望有恰好一个这样的文件,对应于训练的检查点。

您可以直接通过以下方式从Python加载模型

from molecule_generation import load_model_from_directory

model_dir = "./example_model_directory"
example_smiles = ["c1ccccc1", "CNC=O"]

with load_model_from_directory(model_dir) as model:
    embeddings = model.encode(example_smiles)
    print(f"Embedding shape: {embeddings[0].shape}")

    # Decode without a scaffold constraint.
    decoded = model.decode(embeddings)

    # The i-th scaffold will be used when decoding the i-th latent vector.
    decoded_scaffolds = model.decode(embeddings, scaffolds=["CN", "CCC"])

    print(f"Encoded: {example_smiles}")
    print(f"Decoded: {decoded}")
    print(f"Decoded with scaffolds: {decoded_scaffolds}")

这应该会产生一个类似以下输出的结果

Embedding shape: (512,)
Encoded: ['c1ccccc1', 'CNC=O']
Decoded: ['C1=CC=CC=C1', 'CNC=O']
Decoded with scaffolds: ['C1=CC=C(CNC2=CC=CC=C2)C=C1', 'CNC(=O)C(C)C']

如上所示,MoLeR是通过上下文管理器加载的。幕后发生以下操作

  • 首先,选择适当的包装器类:如果提供的目录包含一个MoLeRVae检查点,则返回的包装器将支持encodedecodesample,而MoLeRGenerator将只支持sample
  • 接下来,会启动并行工作进程,这些进程将等待编码/解码的查询;只要上下文活跃,这些进程就会继续存在。可以通过使用num_workers参数来配置并行度。

微调

微调过程与从头开始训练类似,但有几点调整。首先,用于微调的数据必须通过以下命令进行预处理

molecule_generation preprocess INPUT_DIR OUTPUT_DIR TRACE_DIR \
    --pretrained-model-path CHECKPOINT_PATH

其中CHECKPOINT_PATH指向将进行微调的模型的文件(不是目录)。

--pretrained-model-path参数是必要的,否则预处理将仅从提供的SMILES集推断出各种元数据(例如,原子/基团类型集),而对于微调,这必须与模型最初训练时使用的元数据相匹配。

预处理后,微调可以运行如下

molecule_generation train MoLeR TRACE_DIR \
    --load-saved-model CHECKPOINT_PATH \
    --load-weights-only

当在小型数据集上微调时,可能不希望更新模型直到收敛。可以通过传递--model-params-override '{"num_train_steps_between_valid": 100}'(以缩短迷你epoch)和--max-epochs(以限制迷你epoch的数量)来限制训练持续时间。

可视化

我们支持两种细微不同的可视化模式:解码给定的潜在向量,以及解码由编码给定的SMILES字符串创建的潜在向量。在前者的情况下,解码器在推理期间正常运行;在后者的案例中,我们知道真实输入,因此我们通过教师强制来正确解码决策。

要进入可视化器,可以运行以下任一命令

molecule_generation visualise cli MODEL_DIR SMILES_OR_SAMPLES_PATH

以将结果以纯文本形式打印到CLI,或者

molecule_generation visualise html MODEL_DIR SMILES_OR_SAMPLES_PATH OUTPUT_DIR

以将结果保存到OUTPUT_DIR下的静态HTML网页。

代码结构

所有我们的模型都在Tensorflow 2中实现,旨在易于扩展和构建。我们使用tf2-gnn作为核心图神经网络组件。

MoLeR模型本身实现为MoLeRVae类,继承自tf2-gnn中的GraphTaskModel;该基类封装了编码GNN。解码GNN作为一个外部的MoLeRDecoder层实例化;它还包括批量推理代码,该代码在每一步强制选择最大似然。

作者

注意:由于Git历史在开源时被截断,GitHub的统计数据没有反映一些作者的部分贡献程度。上述所有作者都对代码产生了影响,并且(大约)按照贡献度递减的顺序排列。

本代码由英国剑桥微软研究院的生成化学小组维护。我们正在招聘

MoLeR是在我们与诺华研究合作的基础上创建的。特别是,其设计由Nadine SchneiderFinton SirockinNikolaus Stiefl以及诺华的其他同事指导。

贡献

本项目欢迎贡献和建议。大多数贡献需要您同意一份贡献者许可协议(CLA),声明您有权并且实际上已经授予我们使用您贡献的权利。有关详细信息,请访问https://cla.opensource.microsoft.com

当您提交拉取请求时,CLA机器人会自动判断您是否需要提供CLA,并相应地装饰PR(例如,状态检查、注释)。只需遵循机器人提供的说明。您只需在整个使用我们的CLA的仓库中这样做一次。

本项目已采用微软开源行为准则。有关更多信息,请参阅行为准则FAQ或通过opencode@microsoft.com联系以提出任何额外的问题或评论。

风格指南

  • 对于代码风格,请使用blackflake8
  • 对于提交信息,请使用祈使语气,并遵循语义提交信息模板;例如。

    feat(moler_decoder): 改进无效操作的掩码

商标

本项目可能包含项目、产品或服务的商标或徽标。微软商标或徽标的授权使用必须遵循并遵循微软的商标和品牌指南。在此项目的修改版本中使用微软的商标或徽标不得引起混淆或暗示微软的赞助。任何第三方商标或徽标的用途均受这些第三方政策的约束。

项目详情


下载文件

下载适用于您的平台的文件。如果您不确定选择哪个,请了解更多关于安装包的信息。

源代码分发

molecule_generation-0.4.1.tar.gz (957.7 kB 查看哈希值)

上传时间: 源代码

构建分发

molecule_generation-0.4.1-py3-none-any.whl (987.7 kB 查看哈希值)

上传时间: Python 3

支持者

AWSAWS 云计算和安全赞助商 DatadogDatadog 监控 FastlyFastly CDN GoogleGoogle 下载分析 MicrosoftMicrosoft PSF赞助商 PingdomPingdom 监控 SentrySentry 错误日志 StatusPageStatusPage 状态页面