Auto Byte

专注未来出行及智能汽车科技

微信扫一扫获取更多资讯

Science AI

关注人工智能与其他前沿技术、基础学科的交叉研究与融合发展

微信扫一扫获取更多资讯

首个大众可用PyTorch版AlphaFold2复现,哥大开源OpenFold,star量破千

AlphaFold2 是 2021 年 AI for Science 领域最耀眼的一颗星。现在,有人在 PyTorch 中复现了它,并已在 GitHub 上开源。这一复现在性能上媲美原版 AlphaFold2,且在算力、存储方面的要求对于大众来说更加友好。

图片

刚刚,哥伦比亚大学系统生物学助理教授 Mohammed AlQuraishi 在推特上宣布,他们从头训练了一个名为 OpenFold 的模型,该模型是 AlphaFold2 的可训练 PyTorch 复现版本。Mohammed AlQuraishi 还表示,这是第一个大众可用的 AlphaFold2 复现。

图片

AlphaFold2 可以周期性地以原子精度预测蛋白质结构,在技术上利用多序列对齐和深度学习算法设计,并结合关于蛋白质结构的物理和生物学知识提升了预测效果。它实现了 2/3 蛋白质结构预测的卓越成绩并在去年登上了《自然》杂志。更令人惊喜的是,DeepMind 团队不仅开源了模型,还将 AlphaFold2 预测数据做成了免费开放的数据集。

然而,开源并不意味着能用、好用。其实,AlphaFold2 软件系统的部署难度极大,并且对硬件的要求高、数据集下载周期长、占用空间大,每一条都让普通开发者望而却步。因此,开源社区一直在努力实现 AlphaFold2 的可用版本。

这次哥伦比亚大学 Mohammed AlQuraishi 教授等人实现的 OpenFold 总训练时间大约为 100000 A100 小时,但在大约 3000 小时内就达到了 90% 的准确率

图片

OpenFold 与原版 AlphaFold2 的准确率相当,甚至略胜一筹,可能因为 OpenFold 的训练集更大一点:

图片

OpenFold 的主要优势是推理速度显著提升,对于较短的蛋白质序列,OpenFold 的推理速度可以达到 AlphaFold2 的两倍。另外,由于使用自定义的 CUDA 内核,OpenFold 使用更少的内存就能推理更长的蛋白质序列。

OpenFold 介绍

OpenFold 几乎再现了原始开源推理代码 (v2.0.1) 的所有功能,除了已趋于被淘汰的「模型集成」功能,该功能在 DeepMind 自己的消融测试中就表现不佳。

无论是否有 DeepSpeed,OpenFold 都能以全精度或 bfloat16 进行训练。为了实现 AlphaFold2 的原始性能,该团队从头开始训练 OpenFold,现已公开发布了模型权重和训练数据。其中,训练数据包含大约 400000 份 MSA 和 PDB70 模板文件。OpenFold 还支持使用 AlphaFold 的官方参数进行蛋白质推理。

与其他实现相比,OpenFold 具有以下优点:

  • 短序列推理:加快了在 GPU 上推理少于 1500 个氨基酸残基的链的速度;
  • 长序列推理:通过该研究实现的低记忆注意力(low-memory attention)对极长链进行推理,OpenFold 可以在单个 A100 上预测 超过 4000 个残基的序列结构,借助 CPU offload 甚至可以预测更长的序列;
  • 内存高效在训练和推理期间,在 FastFold 内核基础上修改的自定义 CUDA 注意力内核,使用的 GPU 内存分别比等效的 FastFold 和现有的 PyTorch 实现少 4 倍和 5 倍;
  • 高效对齐脚本:该团队使用原始 AlphaFold HHblits/JackHMMER pipeline 或带有 MMseqs2 的 ColabFold,已经生成了数百万个对齐。

Linux 系统下的安装与使用

开发团队提供了一个在本地安装 Miniconda、创建 conda 虚拟环境、安装所有 Python 依赖项并下载有用资源的脚本,包括两组模型参数

运行以下命令:

scripts/install_third_party_dependencies.sh

使用如下命令激活环境:

source scripts/activate_conda_env.sh

停用命令:

source scripts/deactivate_conda_env.sh

在激活环境下,编译 OpenFold 的 CUDA 内核

python3 setup.py install

在 / usr/bin 路径下安装 HH-suite:

# scripts/install_hh_suite.sh

使用如下命令可以下载用于训练 OpenFold 和 AlphaFold数据库

bash scripts/download_data.sh data/

如果要使用一组 DeepMind 的预训练参数对一个或多个序列进行推理,可以运行如下代码:

python3 run_pretrained_openfold.py \
    fasta_dir \
    data/pdb_mmcif/mmcif_files/ \
    --uniref90_database_path data/uniref90/uniref90.fasta \
    --mgnify_database_path data/mgnify/mgy_clusters_2018_12.fa \
    --pdb70_database_path data/pdb70/pdb70 \
    --uniclust30_database_path data/uniclust30/uniclust30_2018_08/uniclust30_2018_08 \
    --output_dir ./ \
    --bfd_database_path data/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
    --model_device "cuda:0" \
    --jackhmmer_binary_path lib/conda/envs/openfold_venv/bin/jackhmmer \
    --hhblits_binary_path lib/conda/envs/openfold_venv/bin/hhblits \
    --hhsearch_binary_path lib/conda/envs/openfold_venv/bin/hhsearch \
    --kalign_binary_path lib/conda/envs/openfold_venv/bin/kalign
    --config_preset "model_1_ptm"
    --openfold_checkpoint_path openfold/resources/openfold_params/finetuning_2_ptm.pt

更多细节请参见 GitHub:https://github.com/aqlaboratory/openfold

图片


扩展阅读:


参考链接:
https://cloud.tencent.com/developer/article/1861192
https://twitter.com/MoAlQuraishi
工程AlphaFold2
相关数据
DeepMind机构

DeepMind是一家英国的人工智能公司。公司创建于2010年,最初名称是DeepMind科技(DeepMind Technologies Limited),在2014年被谷歌收购。在2010年由杰米斯·哈萨比斯,谢恩·列格和穆斯塔法·苏莱曼成立创业公司。继AlphaGo之后,Google DeepMind首席执行官杰米斯·哈萨比斯表示将研究用人工智能与人类玩其他游戏,例如即时战略游戏《星际争霸II》(StarCraft II)。深度AI如果能直接使用在其他各种不同领域,除了未来能玩不同的游戏外,例如自动驾驶、投资顾问、音乐评论、甚至司法判决等等目前需要人脑才能处理的工作,基本上也可以直接使用相同的神经网上去学而习得与人类相同的思考力。

https://deepmind.com/
深度学习技术

深度学习(deep learning)是机器学习的分支,是一种试图使用包含复杂结构或由多重非线性变换构成的多个处理层对数据进行高层抽象的算法。 深度学习是机器学习中一种基于对数据进行表征学习的算法,至今已有数种深度学习框架,如卷积神经网络和深度置信网络和递归神经网络等已被应用在计算机视觉、语音识别、自然语言处理、音频识别与生物信息学等领域并获取了极好的效果。

权重技术

线性模型中特征的系数,或深度网络中的边。训练线性模型的目标是确定每个特征的理想权重。如果权重为 0,则相应的特征对模型来说没有任何贡献。

参数技术

在数学和统计学裡,参数(英语:parameter)是使用通用变量来建立函数和变量之间关系(当这种关系很难用方程来阐述时)的一个数量。

数据库技术

数据库,简而言之可视为电子化的文件柜——存储电子文件的处所,用户可以对文件中的数据运行新增、截取、更新、删除等操作。 所谓“数据库”系以一定方式储存在一起、能予多个用户共享、具有尽可能小的冗余度、与应用程序彼此独立的数据集合。

准确率技术

分类模型的正确预测所占的比例。在多类别分类中,准确率的定义为:正确的预测数/样本总数。 在二元分类中,准确率的定义为:(真正例数+真负例数)/样本总数

AlphaFold技术

DeepMind 提出的深度神经网络蛋白质形态预测方法。AlphaFold系统,是DeepMind在2017-2018年中一直在研究的项目,它建立在多年以前使用大量基因组数据来预测蛋白质结构的研究基础之上。 AlphaFold产生的蛋白质3D模型比以往任何一种都精确得多,在生物学的核心挑战之一上取得了重大进展。

暂无评论
暂无评论~