首個(gè)大眾可用PyTorch版AlphaFold2復(fù)現(xiàn)，哥大開源，star量破千
剛剛，哥倫比亞大學(xué)系統(tǒng)生物學(xué)助理教授 Mohammed AlQuraishi 在推特上宣布，他們從頭訓(xùn)練了一個(gè)名為 OpenFold 的模型，該模型是 AlphaFold2 的可訓(xùn)練 PyTorch 復(fù)現(xiàn)版本。Mohammed AlQuraishi 還表示，這是第一個(gè)大眾可用的 AlphaFold2 復(fù)現(xiàn)。
AlphaFold2 可以周期性地以原子精度預(yù)測(cè)蛋白質(zhì)結(jié)構(gòu)，在技術(shù)上利用多序列對(duì)齊和深度學(xué)習(xí)算法設(shè)計(jì)，并結(jié)合關(guān)于蛋白質(zhì)結(jié)構(gòu)的物理和生物學(xué)知識(shí)提升了預(yù)測(cè)效果。它實(shí)現(xiàn)了 2/3 蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)的卓越成績(jī)并在去年登上了《自然》雜志。更令人驚喜的是，DeepMind 團(tuán)隊(duì)不僅開源了模型，還將 AlphaFold2 預(yù)測(cè)數(shù)據(jù)做成了免費(fèi)開放的數(shù)據(jù)集。
然而，開源并不意味著能用、好用。其實(shí)，AlphaFold2 軟件系統(tǒng)的部署難度極大，并且對(duì)硬件的要求高、數(shù)據(jù)集下載周期長(zhǎng)、占用空間大，每一條都讓普通開發(fā)者望而卻步。因此，開源社區(qū)一直在努力實(shí)現(xiàn) AlphaFold2 的可用版本。
這次哥倫比亞大學(xué) Mohammed AlQuraishi 教授等人實(shí)現(xiàn)的 OpenFold 總訓(xùn)練時(shí)間大約為 100000 A100 小時(shí)，但在大約 3000 小時(shí)內(nèi)就達(dá)到了 90% 的準(zhǔn)確率。
OpenFold 與原版 AlphaFold2 的準(zhǔn)確率相當(dāng)，甚至略勝一籌，可能因?yàn)?OpenFold 的訓(xùn)練集更大一點(diǎn)：
OpenFold 的主要優(yōu)勢(shì)是推理速度顯著提升，對(duì)于較短的蛋白質(zhì)序列，OpenFold 的推理速度可以達(dá)到 AlphaFold2 的兩倍。另外，由于使用自定義的 CUDA 內(nèi)核，OpenFold 使用更少的內(nèi)存就能推理更長(zhǎng)的蛋白質(zhì)序列。
OpenFold 介紹
OpenFold 幾乎再現(xiàn)了原始開源推理代碼 (v2.0.1) 的所有功能，除了已趨于被淘汰的「模型集成」功能，該功能在 DeepMind 自己的消融測(cè)試中就表現(xiàn)不佳。
無(wú)論是否有 DeepSpeed，OpenFold 都能以全精度或 bfloat16 進(jìn)行訓(xùn)練。為了實(shí)現(xiàn) AlphaFold2 的原始性能，該團(tuán)隊(duì)從頭開始訓(xùn)練 OpenFold，現(xiàn)已公開發(fā)布了模型權(quán)重和訓(xùn)練數(shù)據(jù)。其中，訓(xùn)練數(shù)據(jù)包含大約 400000 份 MSA 和 PDB70 模板文件。OpenFold 還支持使用 AlphaFold 的官方參數(shù)進(jìn)行蛋白質(zhì)推理。
與其他實(shí)現(xiàn)相比，OpenFold 具有以下優(yōu)點(diǎn)：
- 短序列推理：加快了在 GPU 上推理少于 1500 個(gè)氨基酸殘基的鏈的速度；
- 長(zhǎng)序列推理：通過(guò)該研究實(shí)現(xiàn)的低記憶注意力（low-memory attention）對(duì)極長(zhǎng)鏈進(jìn)行推理，OpenFold 可以在單個(gè) A100 上預(yù)測(cè) 超過(guò) 4000 個(gè)殘基的序列結(jié)構(gòu)，借助 CPU offload 甚至可以預(yù)測(cè)更長(zhǎng)的序列；
- 內(nèi)存高效在訓(xùn)練和推理期間，在 FastFold 內(nèi)核基礎(chǔ)上修改的自定義 CUDA 注意力內(nèi)核，使用的 GPU 內(nèi)存分別比等效的 FastFold 和現(xiàn)有的 PyTorch 實(shí)現(xiàn)少 4 倍和 5 倍；
- 高效對(duì)齊腳本：該團(tuán)隊(duì)使用原始 AlphaFold HHblits/JackHMMER pipeline 或帶有 MMseqs2 的 ColabFold，已經(jīng)生成了數(shù)百萬(wàn)個(gè)對(duì)齊。
Linux 系統(tǒng)下的安裝與使用
開發(fā)團(tuán)隊(duì)提供了一個(gè)在本地安裝 Miniconda、創(chuàng)建 conda 虛擬環(huán)境、安裝所有 Python 依賴項(xiàng)并下載有用資源的腳本，包括兩組模型參數(shù)。
運(yùn)行以下命令：
scripts/install_third_party_dependencies.sh
使用如下命令激活環(huán)境：
source scripts/activate_conda_env.sh
停用命令：
source scripts/deactivate_conda_env.sh
在激活環(huán)境下，編譯 OpenFold 的 CUDA 內(nèi)核
python3 setup.py install
在 / usr/bin 路徑下安裝 HH-suite：
# scripts/install_hh_suite.sh
使用如下命令可以下載用于訓(xùn)練 OpenFold 和 AlphaFold 的數(shù)據(jù)庫(kù)：
bash scripts/download_data.sh data/
如果要使用一組 DeepMind 的預(yù)訓(xùn)練參數(shù)對(duì)一個(gè)或多個(gè)序列進(jìn)行推理，可以運(yùn)行如下代碼：
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
更多細(xì)節(jié)請(qǐng)參見(jiàn) GitHub：https://github.com/aqlaboratory/openfold
擴(kuò)展閱讀：
- 高效預(yù)測(cè)幾乎所有人類蛋白質(zhì)結(jié)構(gòu)，AlphaFold 再登 Nature，數(shù)據(jù)庫(kù)全部免費(fèi)開放
- 生物計(jì)算專家超細(xì)致解讀 AlphaFold2 論文：模型架構(gòu)及應(yīng)用
- DeepMind 開源的 AlphaFold 怎么用？打開 Colab 就能在線用
