Mila唐建团队发布AI药物研发平台TorchDrug，图灵奖得主Yoshua Bengio等大牛指导

机器之心

发布于 2023-03-29 17:31:49

5390

发布于 2023-03-29 17:31:49

文章被收录于专栏：机器之心

机器之心报道

机器之心编辑部

Mila 唐建团队开源了机器学习药物研发平台 TorchDrug，这是一个专为药物研发设计的机器学习平台。

现代医学的发展使许多曾经的疑难杂症变得不再可怕，但需要正视的是，至今仍有很多种疾病未能找到对症的药物，且药物发现的过程一直是相对漫长和昂贵的。近年来，人工智能在药物研发中的应用受到了越来越多的关注，利用这一领域的技术可以减少人工环节，加速新药研发的速度，缩短药物研发到使用的过程。

近日，Mila 唐建团队开源了机器学习药物研发平台 TorchDrug。TorchDrug 是一个专为药物研发设计的机器学习平台，涵盖从图机器学习（图神经网络、几何深度学习和知识图谱）、深度生成模型到强化学习的技术，它提供了一个全面而灵活的接口来支持 PyTorch 中药物发现模型的快速原型设计。

官网地址：http://torchdrug.ai/

唐建教授表示，药物发现的未来依赖于丰富开源社区的培育。「我们希望这个平台能成为未来机器学习药物发现主要的开源平台，并促进这一方面的进展。」唐建说道。

该平台包含四项关键功能，分别为：

最小化领域知识

面向机器学习社区，TorchDrug 提取了大部分领域知识并提供了基于张量的接口，允许用户使用张量代数和机器学习运算来控制生物医学对象。

数据集与搭建模块

平台拥有大量数据集和搭建模块，用户无需编写示例代码即可在 TorchDrug 中实现标准模型。同时搭建模块也具有高度可扩展性，以促进模型设计的探索。

综合基准测试

研发人员已经对该平台进行了多项药物发现任务的综合基准测试，提供了热门深度学习架构的系统对比。基准测试结果预计将跟随新模型的进步以激发新的研究方向。

可扩展的训练和推理

TorchDrug 具备可扩展的特性，可加速多个 CPU 或 GPU 进行的训练和推理。仅需一行代码，即可让用户在 CPU、GPU 或分布式设置之间无缝切换。

平台功能详解

TorchDrug 是一个基于 PyTorch 的机器学习工具箱，具有如下功能：

基于张量的图

图是 TorchDrug 中的基本数据结构，可以被视为张量框架中图形处理库（例如 NetworkX）的泛化，同时具备 GPU 加速和自动微分的特性。

import torchdrug as tdfrom torchdrug import dataedge_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)

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 批准。

分层接口

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 整个项目由博士生朱兆成、刘圣超、史晨策、张作柏、瞿锰以及多位来自清华大学、北京大学和上海交通大学的实习生共同完成。