Mila唐建团队发布AI药物研发平台TorchDrug，Yoshua Bengio等指导

现代医学的发展使许多曾经的疑难杂症变得不再可怕，但需要正视的是，至今仍有很多种疾病未能找到对症的药物，且药物发现的过程一直是相对漫长和昂贵的。

近年来，人工智能在药物研发中的应用受到了越来越多的关注，利用这一领域的技术可以减少人工环节，加速新药研发的速度，缩短药物研发到使用的过程。

近日，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)

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

同时，本项目也得到了 Yoshua Bengio 等十余位人工智能、生物学领域学者的指导：

此前，唐建等人在 AAAI 2021 中详细介绍了「人工智能与药物研发」的相关内容，感兴趣的读者可以阅读：https://deepgraphlearning.github.io/DrugTutorial_AAAI2021/

并将近期在 KDD2021 中继续介绍「人工智能与药物研发」的相关研究进展：

https://deepgraphlearning.github.io/DrugTutorial_KDD2021/