现代医学的发展使许多曾经的疑难杂症变得不再可怕,但需要正视的是,至今仍有很多种疾病未能找到对症的药物,且药物发现的过程一直是相对漫长和昂贵的。
近年来,人工智能在药物研发中的应用受到了越来越多的关注,利用这一领域的技术可以减少人工环节,加速新药研发的速度,缩短药物研发到使用的过程。
近日,Mila 唐建团队开源了机器学习药物研发平台 TorchDrug。TorchDrug 是一个专为药物研发设计的机器学习平台,涵盖从图机器学习(图神经网络、几何深度学习和知识图谱)、深度生成模型到强化学习的技术,它提供了一个全面而灵活的接口来支持 PyTorch 中药物发现模型的快速原型设计。
官网地址:http://torchdrug.ai/
唐建教授表示,药物发现的未来依赖于丰富开源社区的培育。「我们希望这个平台能成为未来机器学习药物发现主要的开源平台,并促进这一方面的进展。」唐建说道。
该平台包含四项关键功能,分别为:
最小化领域知识
面向机器学习社区,TorchDrug 提取了大部分领域知识并提供了基于张量的接口,允许用户使用张量代数和机器学习运算来控制生物医学对象。
数据集与搭建模块
平台拥有大量数据集和搭建模块,用户无需编写示例代码即可在 TorchDrug 中实现标准模型。同时搭建模块也具有高度可扩展性,以促进模型设计的探索。
综合基准测试
研发人员已经对该平台进行了多项药物发现任务的综合基准测试,提供了热门深度学习架构的系统对比。基准测试结果预计将跟随新模型的进步以激发新的研究方向。
可扩展的训练和推理
TorchDrug 具备可扩展的特性,可加速多个 CPU 或 GPU 进行的训练和推理。仅需一行代码,即可让用户在 CPU、GPU 或分布式设置之间无缝切换。
平台功能详解
TorchDrug 是一个基于 PyTorch 的机器学习工具箱,具有如下功能:
基于张量的图
图
图是 TorchDrug 中的基本数据结构,可以被视为张量框架中图形处理库(例如 NetworkX)的泛化,同时具备 GPU 加速和自动微分的特性。
import torchdrug as td
from torchdrug import data
edge_list = [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [5, 0]]
graph = data.Graph(edge_list, num_node=6)
graph.visualize()
分子
药物发现中的一个图的具体实例是分子,它可能具有额外的化学约束和特征。例如,我们可以从 SMILES 字符串构造一个分子(SMILES 简称分子线性输入规范,是一种用 ASCII 字符串明确描述分子结构的规范)。
mol = data.Molecule.from_smiles("C1=CC=CC=C1")
mol.visualize()
print(mol.node_feature.shape)
print(mol.edge_feature.shape)
torch.Size([6, 69])
torch.Size([12, 19])
图运算
也可以将一组任意图进行批处理,以最大限度地提高硬件利用率。它们也可以使用. cpu() 和. cuda() 方法在 CPU 和 GPU 之间传输。
mols = data.PackedMolecule.from_smiles(["CCSCCSP(=S)(OC)OC", "CCOC(=O)N", "N(Nc1ccccc1)c2ccccc2", "NC(=O)c1cccnc1"])
mols.visualize()
mols = mols.cuda()
print(mols)
![图片](https://image.jiqizhixin.com/uploads/editor/922eea5c-0c07-46a8-8c27-27c5a4887ecd/640.png)
PackedMolecule(batch_size=4, num_nodes=[12, 6, 14, 9], num_edges=[22, 10, 30, 18], device='cuda:0')
就像原始的 PyTorch 张量一样,图支持广泛的索引操作。典型用途包括 node masking, edge masking or graph masking。以下示例显示了如何选择至少包含一个碳的边,更多图形操作可参考注释。
node_in, node_out, bond_type = mols.edge_list.t()
edge_mask = (mols.atom_type[node_in] == td.CARBON) | (mols.atom_type[node_out] == td.CARBON)
mols = mols.edge_mask(edge_mask)
mols.visualize()
分子分类器
TorchDrug 提供了大量热门的用于药物发现和图表示学习的数据集和模型,帮助用户尽可能少地编写和调试示例代码。以下图的 ClinTox 数据集为例,该数据集预测的是一个分子在临床试验中是否有毒,以及它是否被 FDA 批准。
![图片](https://image.jiqizhixin.com/uploads/editor/5763d0ed-3318-4b5d-bfc7-2cd3f71aa939/640.png)
分层接口
TorchDrug 具有分层接口,涵盖从低级数据结构及运算、中间层、模型到其他高级任务。利用低级别的搭建模块,用户也可以轻松地在自定义任何层级的模块。
目前,层次共分为 4 个级别:
torchdrug.data:基本数据结构和基于张量的图运算,e.g. Graph masking;
torchdrug.layers:神经网络层和损失层,e.g. Message Passing Layer;
torchdrug.models:表征学习模型,e.g. Message Passing Neural Network;
torchdrug.tasks:特定任务的模块和功能,e.g. Molecule property prediction。
研发团队
这一平台的项目负责人为加拿大蒙特利尔学习算法研究所(Mila)助理教授唐建,其研究领域包括图表示学习、图神经网络、药物发现及知识图谱。唐建 2014 年毕业于北京大学信息科学技术学院并获得博士学位,2014-2016 年任职微软亚洲研究院副研究员,2016-2017 年为密歇根大学和卡内基梅隆大学的联合培养博士后,曾获 2014 年机器学习顶级会议 ICML 的最佳论文。唐建是图表示学习领域的代表性学者,他所提出的网络表示学习算法 LINE 被广泛应用,其他代表工作还包括 RotatE 等。
Mila 实验室是由深度学习先驱 Yoshua Bengio 教授领导的人工智能实验室(https://mila.quebec/),主要从事深度学习、强化学习、优化算法等人工智能领域的基础研究以及在不同领域的应用。TorchDrug 整个项目由博士生朱兆成、刘圣超、史晨策、张作柏、瞿锰以及多位来自清华大学、北京大学和上海交通大学的实习生共同完成。
此前,唐建等人在 AAAI 2021 中详细介绍了「人工智能与药物研发」的相关内容,感兴趣的读者可以阅读:https://deepgraphlearning.github.io/DrugTutorial_AAAI2021/
并将近期在 KDD2021 中继续介绍「人工智能与药物研发」的相关研究进展:
https://deepgraphlearning.github.io/DrugTutorial_KDD2021/