萤火跑模型 | Alphafold 蛋白质结构预测

Vachel November 05, 2021

2021年刚刚走过，要说去年一年AI领域的最重大进展，DeepMind在Nature上发表的“AlphaFold”无出其右，其旨在根据一个蛋白质的氨基酸序列来确定它的3D结构，是生物学领域中的几个重大挑战课题之一。在“蛋白质结构预测关键评估”（CASP）的蛋白质结构预测年赛上，AlphaFold击败了其他百来支团队并取得了惊人的效果，有些时候，AlphaFold预测的结构与利用X射线晶体学和近几年的冷冻电镜（cryo-EM）等“金标准”实验方法所确定的结构几乎一样。科学家说，目前看来，AlphaFold暂时还不能取代这些费力又昂贵的实验方法，但它将带来全新的研究生命的方式。

作为一个计算机专业的算法从业人员，本期笔者带着浓厚的兴趣来尝试一下 DeepMind 开源的 Alphafold 模型并验证效果，试跑的平台是幻方AI团队研发的“萤火二号”。我们复现了 Alphafold 的训练与预测过程，为大家带来第一手的测试体验。

论文标题：Improved protein structure prediction using potentials from deep learning

原文地址：https://www.nature.com/articles/s41586-019-1923-7

源码地址：https://github.com/aqlaboratory/openfold

模型仓库：https://github.com/HFAiLab/alphafold-optimized

模型介绍

AlphaFold 解决的问题是蛋白质折叠问题。输入是一个氨基酸序列，每一个位置代表一个元素，输出是一个拓扑结构，如下图所示：

直接输出一个拓扑结构对于深度学习来说比较困难，一般需要输出结构化的数据。DeepMind 将该问题转化为预测氨基酸链的一些性质。准确的来说，深度模型输出的是每一个氨基酸单元和其下一个氨基酸单元在空间中的夹角与距离，再组合成拓扑结构。

如上图所示，AlphaFold算法大致分为以下几个部分：

特征工程：序列和MSA特征抽取，结合专家经验数据库，把氨基酸链的输入转换到特征空间；
深度神经网络结构预测：依据1中的特征预测氨基酸链的一些性质，比如氨基酸之间两两的距离分布，氨基酸链的夹角分布；
Potential Construction：结合专家经验构造一个评估函数，来评估2中神经网络输出解的合理程度；
结构生成：对于2中预测的距离分布、夹角分布，使用3中的评估函数评估Loss，然后使用梯度下降法优化，直到收敛。

可以看到，Alphafold 核心贡献在于：

深度学习模型融入专家经验：特征构建阶段，尽可能结合领域知识构建特征，以提高模型表现；评估阶段，将专家知识（氨基酸链的物理特性）结合到损失函数中，进行更精准的判断。
全新的蛋白质折叠问题的建模方式：DeepMind不是直接去预测蛋白质的结构，而是将该问题转化为预测氨基酸折叠后的一些性质，再将这些性质转化为可微的约束，并求解。这种思路可以应用在很多复杂的求解问题上。
对于梯度下降算法的优化：DeepMind提出通过类似进化/退火算法，解决梯度下降中的局部最优问题，以求解真正的“最优解”。

模型实践

DeepMind 仅开源了 Alphafold 的预测代码和预训练好的模型参数，而对于训练部分，OpenFold 与 UniFold 分别开源了 Pytorch 和 Tensorflow 的版本。我们选择 OpenFold 进行改造，因为其复现完整度比较高，且使用 Pytorch 框架更适配萤火集群的运行环境。项目地址： https://github.com/aqlaboratory/openfold

1. 数据集

从上文可以知道，Alphafold 数据集包含两个部分：氨基酸序列数据、专家经验数据。具体的，数据说明如下：

bfd / small_bfd：Big Fantastic Database，由论文作者 Martin Steinegger 构建的蛋白质序列数据库，包含65,983,866个蛋白质家族（MSAs）和2,204,359,010个蛋白质序列（HMMs），总计7个文件1.8T数据量。
mgnify：一个微生物基因组数据库，使用JackerHammer工具进行检索，总计一个文件64G数据量。
pdb70：一个蛋白质结构数据库，总计9个文件56G数据量。
pdb_mmcif：一个蛋白质结构数据库，包含超过55,000种不同蛋白质的180,000多个3D结构，总计约18万个文件206G数据量。
uniclust30：1.3亿条氨基酸序列，使用hhsuite中的HHblits工具进行快速msa检索，总计13个文件86G数据量。
uniref90：1.3亿条氨基酸序列，是msa检索所需要的库，使用JackerHammer工具进行检索，总计1个文件58G数据量。

在萤火集群上，这些数据如果以独立文件形式直接用于训练，小文件的高频读写会给集群的存储系统造成不小压力。对应的，我们需要将数据转化成 FFRecord 数据格式，合并多个小文件以减小磁盘 IO 的开销。FFRecord (FireFlyer Record) 数据格式是幻方AI自研的简单高效的存储二进制记录的文件格式。为了方便研究者的使用，我们将以上的多个数据集进行整合，加入hfai 高性能数据集仓库，对外提供 hfai.datasets.AlphafoldData 接口。具体的使用案例如下：

from hfai.datasets import AlphafoldData

# 将数据转为Alphafold训练的格式
transform = AlphafoldDataTransform()
dataset = AlphafoldData(transform)
dataloader = dataset.loader(batch_size=1, num_worker=8)

for data in dataloader:
    # training iter

AlphafoldData 数据集接口返回的数据是一个元组，其中包含了蛋白质本身的信息 pdb_code， mmcif_string 和在蛋白质数据库中匹配得到的结果 bfd_hits，mgnify_hits，pdb70_hits，uniref90_hits。这些数据在训练中需要被 AlphafoldDataTransform 进一步处理，具体的使用案例可以访问我们开源在模型仓库中的 Alphafold-optimzied。

除了以上整合的数据，pdb_mmcif 数据在 Alphafold 训练后期指导着模型生成合适的蛋白质结构，其具有将近 18w 个小文件，同样需要将其转化成 FFRecord 数据格式。

原始数据转 FFRecord：

from ffrecord import FileWriter

out_file = os.path.join(mmcif_path, "mmcifs.ffr")
writer = FileWriter(outfile, 184929)
num, mapids = 0, {}
for file in os.listdir(os.path.join(mmcif_path, "mmcif_files")):
    with open(os.path.join(mmcif_path, "mmcif_files", file), "r") as fp:
        mmcif_string = fp.read()
        ms = pickle.dumps(mmcif_string)
        writer.write_one(ms)
        mapids[file[:-4]] = num
        num += 1
with open(os.path.join(mmcif_path, "mmcifs_index.pk"), "wb") as fp:
    pickle.dump(mapids, fp)

读取 FFRecord 文件：

from ffrecord import FileReader

mmcif_reader = FileReader(os.path.join(mmcif_dir, "mmcifs.ffr"), check_data=True)
with open(os.path.join(mmcif_path, "mmcifs_index.pk"), "rb") as fp:
    mmcif_index = pickle.load(fp)
if hit_pdb_code in mmcif_index:
    idx = [mmcif_index[hit_pdb_code]]
    data = mmcif_reader.read(idx)
    cif_string = pickle.loads(data[0])

以上数据皆已集成进hfai 数据集仓库，如果需要利用该数据集进行科研的朋友可以直接拉取。

2. 预处理

从上文我们知道，Alphafold 需要与专家经验库进行交互，抽取 MSA 特征。这里用到了多个生物学工具：

openmm：用于分析模拟的高性能工具包，是 Omnia（预测生物分子模拟的工具套件）的一部分。
kalign：一款针对大规模基因组序列的多序列对比工具，运行速度与准确度都比较高。
jackhmmer：一款高效的蛋白质搜索工具，用于从序列数据库中，找到同源的序列。
pdbfixer：用于修复PDB数据，对其中蛋白质结构进行基本的准备或者修复。
HHsuite：一款蛋白质研究软件，快速搜索分析蛋白质的功能特性。

以上工具都只能使用CPU多核进行加速计算，对CPU的算力比较依赖。对应的，在幻方萤火的系统结构中，GPU算力池与CPU算力池是分离开的。因此，数据预处理与模型计算要拆分成两个模块单独进行。

name: precompute_alignments
code:
	directory: /ceph-jd/prod/jupyter/hwj/notebooks/openfold
  entrypoint: scripts/precompute_alignmnets.py
  parameters: mmcif/ alignments/ --cpus 250
resource:
	- container: cuda_11
  - node_count: 1
    group: jd_dev_alpha
priority: 20

这里，jd_dev_alpha 是CPU计算池。通过 hfai run提交上述yaml 文件便可以发起对于目标蛋白质结构的预处理任务。

3. 模型训练

openfold 采用了 Pytorch_lightning 来构建整个训练代码，其自动封装了 dataloader、多机多卡训练、checkpoint 保存与打断自恢复等诸多逻辑，可以直接适配萤火二号的智能分时调度系统。这里我们通过 hfai run运行下述Yaml文件：

name: precompute_alignments
code:
	directory: /ceph-jd/prod/jupyter/hwj/notebooks/openfold
  entrypoint: train_openfold.py
  parameters: mmcif/ alignments/ data/pdb_mmcif output 2021-10-10 --gpus 1 --resume_from_ckpt output/resume_checkpoint.ckpt
resource:
	- container: cuda_11
  - node_count: 1
    group: jd_a100
priority: 20

这里，jd_a100是GPU计算池。由于目标预测的样本不多，这里使用一个计算卡进行训练，过程如下图所示：

训练结束后，我们加载输出的训练模型，预测蛋白质的结构，得到PDB文件。通过PyMol3d 工具，可以将PDB中的蛋白质结构数据可视化出来，最终效果如下：

体验总结

Alphafold作为今年火爆全球的人工智能领域重点突破之一，开启了AI+生物医药领域的新时代。借助萤火二号和自研文件系统，我们可以比较高效地训练该模型。同时，该模型不同于以往深度学习训练场景，其特征抽取与训练，对CPU与GPU的计算需求都比较高。在我们的测试中，GPU整体使用率不高，在20%~30%左右，训练的时间瓶颈是在CPU的计算上。我们欢迎感兴趣的学者们加入我们，继续在幻方AI团队的萤火平台上对该课题进行进一步探索。综合体验打分如下：

研究指数：★★★★★

该模型开创了AI+生物的新课题，引领了一个新的AI时代。
开源指数：★★★

数据已开源，代码开源了预测部分，训练代码开源需要依赖第三方实现。
门槛指数：★★★★★

数据量大，需要CPU与GPU双重算力支持，对算力的需求比较大。
通用指数：★★

模型专注于蛋白质结构预测，有大量的专家经验结合，通用性比较差。
适配指数：★★

需要CPU与GPU的双重算力，整体GPU利用率不高。同时代码和数据都有一定改造工作。

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力AI落地与价值创造，欢迎各方数据研究者与开发者们一同共建。