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