今天,为了加快生物制药领域的研究,从创造疾病的治疗方法到生产新的合成生物材料,我们宣布了一个新的顶点人工智能解决方案,展示了如何使用顶点人工智能管道来大规模地运行DeepMind的AlphaFold蛋白质结构预测。
一旦确定了蛋白质的结构,了解了它在细胞中的作用,科学家就可以根据蛋白质在细胞中的作用开发出能够调节其功能的药物。DeepMind是Alphabet旗下的一个人工智能研究机构,它创建了AlphaFold系统,通过帮助数据科学家和其他研究人员准确地大规模预测蛋白质的几何结构来推动这一领域的研究。
2020年,在 "蛋白质结构预测技术关键评估"(CASP14)实验中,DeepMind展示了AlphaFold的一个版本,对蛋白质结构的预测非常准确,专家们宣布"蛋白质折叠问题"已经解决。第二年,DeepMind开放了AlphaFold 2.0系统的源代码。不久之后,谷歌云发布了一个解决方案,将AlphaFold与顶点人工智能工作台整合,以促进互动实验。这使得许多数据科学家更容易有效地使用AlphaFold,而今天的公告是建立在这个基础上的。
上周,DeepMind与欧洲生物信息学研究所(EMBL-EBI)合作,发布了 科学界已知的几乎所有编目蛋白质的预测结构,AlphaFold又向前迈出了重要一步。这次发布将AlphaFold数据库从近100万个结构扩展到2亿多个结构--这可能会在很大程度上增加我们对生物学的理解。在AlphaFold数据库的这种持续增长和顶点人工智能的效率之间,我们期待着世界各地的研究人员能够有所发现。
在这篇文章中,我们将解释如何开始实验这个解决方案,我们还将调查它的好处,其中包括通过优化选择硬件提供更低的成本,通过实验跟踪、血统和元数据管理提供可重复性,以及通过并行化提供更快的运行时间。
在顶点AI上运行AlphaFold的背景
生成一个蛋白质结构预测是一个计算密集型的任务。它需要大量的CPU和ML加速器资源,可能需要几个小时甚至几天的时间来计算。规模化运行推理工作流程可能具有挑战性--这些挑战包括优化推理时间,优化硬件资源利用率,以及管理实验。我们新的顶点人工智能解决方案旨在解决这些挑战。
为了更好地了解该解决方案如何应对这些挑战,让我们回顾一下AlphaFold的推理工作流程:
-
特征预处理:你使用输入的蛋白质序列(FASTA格式),使用常见的开源工具在跨生物体的遗传序列和蛋白质模板数据库中搜索。这些工具包括JackHMMER与MGnify和UniRef90,HHBlits与Uniclust30和BFD,以及HHSearch与PDB70。搜索的结果(包括多序列比对(MSA)和结构模板)和输入序列被处理为推理模型的输入。你可以只在CPU平台上运行特征预处理步骤。如果你使用全尺寸的数据库,这个过程可能需要几个小时来完成。
-
模型推理:AlphaFold结构预测系统包括一组预训练的模型,包括预测单体结构的模型,预测多聚体结构的模型,以及为CASP微调过的模型。在推理时,你在同一组输入上独立运行一个特定类型的五个模型(如单体模型)。默认情况下,折叠单体模型时每个模型生成一个预测,折叠多聚体时每个模型生成五个预测。推理工作流程的这一步在计算上非常密集,需要GPU或TPU加速。
-
(可选)结构松弛:为了解决推理模型返回的结构中存在的任何结构违规和冲突,你可以执行一个结构松弛步骤。在AlphaFold系统中,你可以使用OpenMM分子力学模拟包来执行一个受限的能量最小化程序。松弛也是非常密集的计算,虽然你可以在一个纯CPU的平台上运行这个步骤,但你也可以通过使用GPU来加速这个过程。
顶点AI解决方案
采用顶点AI解决方案的AlphaFold批处理推理,可以让你通过专注于以下优化,有效地运行AlphaFold推理的规模:
-
通过独立步骤的并行化来优化推理工作流程。
-
通过在最佳硬件平台上运行每个步骤,优化硬件利用率(以及因此而产生的成本)。作为这种优化的一部分,该解决方案自动提供和删除一个步骤所需的计算资源。
-
描述了一个强大而灵活的实验跟踪方法,简化了运行和分析数百个并发推理工作流的过程。
下图显示了该解决方案的架构:
该解决方案包含以下内容:
-
管理遗传数据库的策略:该解决方案包括高性能、完全管理的文件存储。在这个解决方案中,Cloud Filestore被用来管理多个版本的数据库,并提供高吞吐量和低延迟的访问。
-
一个协调器来并行化、协调和有效地运行工作流程中的步骤:预测、松弛和一些特征工程可以被并行化。在这个解决方案中,Vertex AI Pipelines被用作工作流程步骤的协调器和运行时执行引擎。
-
为每个步骤选择优化的硬件平台:预测和放松步骤在GPU上运行,而特征工程在CPU上运行。预测和放松步骤可以使用多GPU节点的配置。这对预测步骤尤其重要,因为内存的使用量与残基的数量呈近似的二次方。因此,预测一个大的蛋白质结构可以超过单个GPU设备的内存。
-
元数据和人工制品管理:该解决方案包括对大规模运行和分析实验的管理。在这个解决方案中,顶点AI元数据被用来管理元数据和人工制品。
该解决方案的基础是一套可重复使用的顶点AI管道组件,它封装了AlphaFold推理工作流程的核心步骤:特征预处理、预测和放松。除了这些组件外,还有一些辅助组件将特征工程步骤分解成工具,以及帮助组织和协调工作流程的辅助组件。
该解决方案包括两个样本管道:通用管道和单体管道。通用管道反映了AlphaFold Github资源库中推理脚本的设置和功能。它跟踪经过的时间并优化计算资源的利用。单体管道进一步优化了工作流程,使特征工程更加高效。你可以通过插入你自己的数据库来定制该管道。
接下来的步骤
要了解更多并尝试这个解决方案,请查看我们的GitHub仓库,其中包含组件和通用及单体管道。仓库中的工件是为了让你可以定制它们。此外,你可以将这个解决方案整合到你的上游和下游工作流程中,以便进一步分析。要了解更多关于顶点人工智能的信息,请访问我们的产品页面。
