Mila唐建团队开源大分子机器学习平台TorchProtein：分析蛋白质序列及结构数据，仅需一两行代码

继药物研发机器学习平台 TorchDrug 之后，时隔一年，Mila 唐建团队开源了新的蛋白质机器学习平台 TorchProtein，这是目前第一个专门针对蛋白质研究的开源机器学习库。

蛋白质是生物体的重要组成成分。理解蛋白质的结构与生化性质，对于药物研发和人类健康有着不可估量的意义。传统基于生物实验的蛋白质研究不仅周期漫长，而且开销巨大。相比之下，机器学习技术则能大幅降低蛋白质研究的周期和开销，为新药的研发带来革命性的影响。然而，基于机器学习的蛋白质研究，涉及到生物领域知识、机器学习算法、并行实现等多个方面，具有较高的入门门槛。市面上也缺少合适的开源库来支持这方面的研究，致使机器学习技术在蛋白质研究中发展受阻。

近日，Mila 唐建团队联合英伟达、英特尔、IBM 以及蛋白质设计初创公司百奥几何共同开源了蛋白质机器学习平台 TorchProtein。TorchProtein 在此前开源平台 TorchDrug 的基础上，为蛋白质打造了一套专用的模块组件。TorchProtein 不仅提供了处理蛋白质的数据结构、主流的算法模型，还包括了标准数据集和任务评测接口。其所有接口均有很强的可扩展性，满足各类机器学习算法开发的需要。无论是图机器学习、蛋白质语言模型还是自监督训练，都能轻松基于 TorchProtein 实现。

官网教程：https://torchprotein.ai/tutorials
基于 TorchProtein 的相关研究：
GearNet（https://arxiv.org/abs/2203.06125）
PEER Benchmark （https://arxiv.org/abs/2206.02096）

唐建教授表示，未来机器学习辅助下的蛋白质研发依赖于丰富开源社区的培养。「我们期待这个平台能成为未来机器学习蛋白质研发主要的开源平台，并推动这一方面的进展。」唐建说道。

TorchProtein 平台的四大核心优势：

统一分子、序列和结构信息的数据结构

考虑到不同任务需要用到蛋白质分子、序列和结构等不同信息，TorchProtein 设计了一套统一不同模态信息的数据结构。可对蛋白质进行分子、序列或结构层面的操作，并在模态之间进行无缝切换。

灵活的算法构建模块

平台提供了多种基于蛋白质序列与结构的模型，仅需一两行代码即可调用 TorchProtein 中的标准模型来分析蛋白质序列和结构数据。针对在蛋白质结构上构图繁琐的问题，TorchProtein 还专门提供了灵活的即时构图的模块，支持在 GPU 上动态构图。

大量基准测试结果

TorchProtein 中引入了大量蛋白质数据集和相关基准测试任务，并记录了主流的机器学习算法在这些测试任务上的测试结果，为新的算法研究提供代码与实验支持。

蛋白质预训练模型

针对基于蛋白质序列和结构的预测任务，平台提供了许多大规模预训练模型，这些模型将有效促进蛋白质机器学习在实际中的运用，并大大缩减计算成本。

平台功能详解

下面我们将从数据结构到算法模型，依次介绍 TorchProtein 的功能：

蛋白质数据结构

统一的蛋白质数据结构

TorchProtein 使用统一的图数据结构来表征蛋白质序列和结构，是对 TorchDrug 中图数据结构的特化。我们可以通过蛋白质结构 PDB 文件来构建数据，并指定使用原子、氨基酸残基和化学键特征。

from torchdrug import data, utilsfrom rdkit import Chemimport nglview
pdb_file = utils.download("https://files.rcsb.org/download/2LWZ.pdb", "./")mol = Chem.MolFromPDBFile("2LWZ.pdb")view = nglview.show_rdkit(mol)view


    

protein = data.Protein.from_pdb(pdb_file, atom_feature="position", bond_feature="length", residue_feature="symbol")print(protein)

Protein(num_atom=445, num_bond=916, num_residue=57)

利用这一数据结构，我们可以轻松获取蛋白质序列信息，TorchProtein 同样支持从蛋白质序列来构建数据结构。

aa_seq = protein.to_sequence()print(aa_seq)seq_protein = data.Protein.from_sequence(aa_seq, atom_feature="symbol", bond_feature="length", residue_feature="symbol")print(seq_protein)

FVNQHLCGSDLVEALYLVCGERGFFYTDPTGGGPRRGIVEQCCHSICSLYQLENYCNProtein(num_atom=445, num_bond=910, num_residue=57)

蛋白质操作

为了充分利用 GPU 资源，TorchProtein 支持将多个蛋白质打包处理，且数据可以在 CPU 和 GPU 之间自由切换。

proteins = [protein] * 3proteins = data.Protein.pack(proteins)proteins = proteins.cuda()print(proteins)

PackedProtein(batch_size=3, num_atoms=[445, 445, 445], num_bonds=[916, 916, 916], num_residues=[57, 57, 57], device='cuda:0')

与 PyTorch 中的张量类似，TorchProtein 的蛋白质数据结构支持按照氨基酸残基的位置进行索引、切分和重组。

segments = [protein[:2], protein[2:4], protein[4:6], protein[6:8]]segments = data.Protein.pack(segments)segments.visualize()

TorchProtein 也提供了动态建图、蛋白质序列和结构分割、原子级别和氨基酸级别结构切换等功能，详细操作指南可参考官网教程。

蛋白质序列性质分析

TorchProtein 提供了许多用于蛋白质序列性质分析的数据集、任务和模型，尽可能避免重复编写代码，帮助用户快速评估各模型在各数据集上的性能。

以蛋白质在细胞中位置预测数据集（Subcellular Localization）为例，我们可以通过两行代码构建数据集并获取其训练集、验证集和测试集。

from torchdrug import datasets
dataset = datasets.SubcellularLocalization("~/protein-datasets/", residue_only=True)train_set, valid_set, test_set = dataset.split()

接着，我们可以定义一个简单的两层 CNN 模型作为蛋白质序列编码器，并在这个 CNN 基础上定义任务模型来进行蛋白质细胞位置预测。

from torchdrug import core, models, tasks
model = models.ProteinCNN(input_dim=21, hidden_dims=[1024, 1024],                        kernel_size=5, padding=2, readout="max")task = tasks.PropertyPrediction(model, task=dataset.tasks, criterion="ce", metric=("acc", "mcc"), num_mlp_layer=2)

TorchProtein 提供了多功能的求解器（solver）来执行模型训练、评测和模型参数保存。

import torch
optimizer = torch.optim.Adam(task.parameters(), lr=1e-4)solver = core.Engine(task, train_set, valid_set, test_set,optimizer, batch_size=64, gpus=[0])solver.train(num_epoch=100)solver.evaluate("valid")solver.save("subloc_cnn.pth")

在 TorchProtein 平台上，研发团队评测了各蛋白质序列机器学习模型在各基准任务上性能，在蛋白质细胞位置预测基准上的结果如下。详细基准测试结果见 PEER Benchmark 论文（https://arxiv.org/abs/2206.02096）。

蛋白质结构性质分析

TorchProtein 也提供了多种用于蛋白质结构性质分析的数据集和模型，以供用户进行模型评估，进而促进蛋白质结构分析的实际应用。

以 Enzyme Commission（EC）蛋白质功能预测数据集为例，通过两行代码我们可以轻松构建数据集并获取其训练集、验证集和测试集。

dataset = datasets.EnzymeCommission("~/protein-datasets/")train_set, valid_set, test_set = dataset.split()

TorchProtein 提供了各种蛋白质结构编码器的实现，这里我们选择使用当前最优之一的 GearNet-Edge 作为蛋白质结构编码器，并在此基础上构建任务模型以解决 EC 数据集中多个二分类问题。




    
model = models.GearNet(input_dim=21, hidden_dims=[512, 512, 512, 512, 512, 512],                        num_relation=7, edge_input_dim=59, num_angle_bin=8,                    batch_norm=True, concat_hidden=True, short_cut=True, readout="sum")task = tasks.MultipleBinaryClassification(model, num_mlp_layer=3, criterion="bce",task=[_ for _ in range(len(dataset.tasks))], metric=["auprc@micro", "f1_max"])

同样，我们可以利用 TorchProtein 的多功能求解器来进行模型训练、评测和模型参数保存。

import torch
optimizer = torch.optim.Adam(task.parameters(), lr=1e-4)solver = core.Engine(task, train_set, valid_set, test_set,optimizer, batch_size=4, gpus=[0])solver.train(num_epoch=100)solver.evaluate("valid")solver.save("ec_gearnet_edge.pth")

TorchProtein 团队在 EC 和 GO 两个蛋白质功能预测基准上评测了各种蛋白质结构编码模型的性能，结果如下。实现和评测细节见 GearNet 论文（https://arxiv.org/abs/2203.06125）。

开发团队

这一平台的项目负责人为加拿大蒙特利尔学习算法研究所（Mila）副教授、终身教授唐建，其研究领域包括几何深度学习、图表示学习、图神经网络、药物发现及知识图谱。唐建 2014 年毕业于北京大学信息科学技术学院并获得博士学位，2014-2016 年任职微软亚洲研究院副研究员，2016-2017 年为密歇根大学和卡内基梅隆大学的联合培养博士后，曾获 2014 年机器学习顶级会议 ICML 的最佳论文。唐建是图表示学习领域的代表性学者，他所提出的网络表示学习算法 LINE 被广泛应用，其他代表工作还包括 RotatE 等。

Mila 实验室是由深度学习先驱 Yoshua Bengio 教授领导的人工智能实验室（https://mila.quebec/），主要从事深度学习、强化学习、优化算法等人工智能领域的基础研究以及在不同领域的应用。TorchProtein 整个项目由博士生张作柏、徐明皓、朱兆成、袁新钰，以及多位来自蛋白质设计初创公司百奥几何、IBM 研究院、英特尔和英伟达的工业界合作者，以及多位来自剑桥大学、清华大学和北京大学的实习生共同完成。