技術乾貨 | AlphaFold/ RoseTTAFold開源復現(1)—推理復現

最近AlphaFold開源，比較火，項目組也嘗試進行復現，有些經驗給大家分享一下，包括推理復現、訓練復現、分布式訓練復現等，今天先介紹一下推理的復現。

 01 

蛋白質幾乎參與所有生命現象，從催化化學反應的酶、到對抗病毒的抗體、以及作為信號物質的胰島素。決定蛋白質功能的是由胺基酸序列摺疊形成三維結構，各種蛋白質相互結合從而去影響生命現象。因此有「序列決定結構，結構決定功能」的說法。

1972年諾貝爾化學獎的獲獎演說中，克裡斯蒂安·安芬森（Christian Anfinsen）提出了一個著名的假設：理論上，蛋白質的胺基酸序列應該完全決定它的結構。這一假設引發了一個長達50年的探索，即能夠僅根據蛋白質的一級胺基酸序列來計算預測蛋白質的三維結構。然而，一個主要的挑戰是，理論上一種蛋白質在形成最終的三維結構之前可以摺疊的方式是天文數字。

基於高通量的測序技術的發展，使用DNA測序，從而測定蛋白質序列相對比較快速和便宜；然而，通過實驗解析蛋白質結構則耗時長，成本高，難度大。常用的方法有X線晶體衍射圖譜法，核磁共振法，冷凍電鏡法等。 

在UniProt資料庫中，有超過2億條以上的序列數據；然而在PDB數據中，僅有超過17萬條以上（有部分重複）的結構數據。大部分的蛋白序列都沒有結構數據。很自然的，如果能夠通過序列預測結構，是很有科學價值的事情。

1994年，John Moult教授和Krzysztof Fidelis教授創立了CASP比賽。在CASP13以前，主要是以各種傳統非端到端的預測的方法為主的,學術界也開始嘗試引入深度學習，但離實驗測得的結構差距較遠。DeepMind在CASP13提出的AlphaFold1取得了顯著的進步，CASP14參賽的全新實現的AlphaFold2更是在很多蛋白質的預測上接近實驗數據。

為了後面的分享，我們再補充一些最必要的知識：

蛋白質的主鏈Backbone和側鏈Sidechain構成。蛋白質的主鏈指的是那條肽鏈；側鏈是每個胺基酸上決定不同胺基酸的部分（R）（胺基酸四個取代基NH2 ,COOH, H + 側鏈R）。

（說明：所有圖片，來自於網絡, 下同。）

胺基酸有20種，不同的胺基酸構成的序列，則為一級結構；部分子序列形成了相對固定的二級結構，如alpha-helix螺旋, beta-sheet帶子；最終摺疊成三級結構以及四級結構

（註：圖中的結構，為PyMOL等可視化軟體卡通化的結果。）

AlphaFold，CASP比賽，還有解蛋白質結構的生物學家，最終要解決就是通過蛋白質的胺基酸序列解析其3D結構。結構決定功能，最終在疾病治療，生物製藥等方面發揮作用。

為了理解算法，這些概念也比較重要：

Residue，殘基，在蛋白質的序列中，胺基酸之間的氨基和羧基脫水成鍵，胺基酸由於其部分基團參與了肽鍵的形成，剩餘的結構部分則稱胺基酸殘基。

Co-Evolution，共同進化，意思就是兩個相關蛋白在進化過程中，一者發生突變，另一個蛋白對應的位點通常也會發生突變以繼續保證二者之間的關聯。做預測結構這個挺常用的，用來定位距離。

主鏈-側鏈和各種角：

簡單的說，這個問題就是：給定蛋白質的序列數據，預測出其3D結構。

其3D結構的表示，可以預測原子間距離，也可預測距離和角度，還可以直接預測原子的坐標等不同的做法；也可以先預測出中間結果殘基的接觸關係然後再預測出最終結果。

如果僅僅通過17萬有結構的數據來直接訓練，效果比較差。一般都會基於輸入的序列，去尋找共進化的序列或序列的Embedding為原始輸入的補充，這個過程即MSA多序列對齊；另外還會去搜索輸入序列的同源的蛋白，這個過程即Template模板搜索。除了AlphaFold團隊，比較知名的研究團隊還有華盛頓大學的David Baker教授，密西根大學張陽教授，芝加哥豐田計算技術研究所許錦波教授等。這次Science發表的RoseTTAFold算法的David團隊，在算法中還引入了二級結構的搜索結果。

（說明：關於蛋白質相關生物問題描述，如有錯誤，在所難免。）

 02 

AlphaFold的開源，在論文、補充材料和代碼中，給出了非常詳細的信息，除了沒有訓練用的代碼和訓練數據（如MSA和Template數據，需要用相關工具從原始數據上生成）。

上圖輸入側的四條線，分別表示：原始的輸入序列；通過hhblits/JackHMMER搜索的MSA（序列-殘基）；搜索的結構和序列產生的Pair（殘基-殘基）；通過HHsearch搜索的Template。

在算法部分，為主幹為EvoFormer，結構部分為等變Transformer，且在兩部分做「迭代」提升。DeepMind認為，EvoFormer核心是MSA 內交換信息的新機制，並且其配對表示允許直接推理空間和進化關係，而Equivariant Attention主要是能適應每個殘基的旋轉和平移，更精細的優化主鏈側鏈裡的結構。

在CASP14後基於DeepMind透露的有限的信息，大家在嘗試復現AlphaFold2的各種技術點，對比這些技術點：

算法中廣泛使Attention，各種Attention變體如Axial Attention，引入等變Transformer，Pairwise Representation，Graph表達（殘基作為節點，殘基間關係作為邊），端到端直接預測坐標或等價的表示，充分利用MSA和Template的共進化、同源等信息等；

RoseTTAFold也包含逐步優化的思想（在Baker的RoseTTAFold算法中，近似的迭代優化為Iterative Feature Extractor，其迭代塊為Transformer,ResNet,SE3。）

開源的AlphaFold2的聯合從Label和Unlabeled數據中學習這部分，是復現工作沒有涉及的。其做法為：先用17萬中有X序列-Y坐標的先訓練，然後用初訓練的網絡對Uniclust30更多的（350,000）只有X序列的做預測和評估，對預測的好的（high-confidence subset）也作為訓練樣本，然後做數據增強（cropping and MSA subsampling），最終取得更好的效果，該技術點為noisy student self-distillation，在很多其他場合被使用。

 03 AlphaFold2/RoseTTAFold算法復現

AlphaFold的訓練部分的復現，需要幾部分的必要工作：

1）大數據集的下載（數百G）和MSA/Template的生成，生成非常耗時，這個需要大量CPU伺服器並行，生成的數據在數T以上；以Sequence長度PDB資料庫中50-60的短序列的MSA生成單伺服器，hhblits預設參數，都在1周左右。 

2）基於jax，haiku的訓練代碼的實現和並行訓練。

3）初訓練和擴展數據集的形成。

後面兩塊AlphaFold2還沒有看到看到開源實現。

如果您對從零完整復現訓練感興趣，可以關注我們後續復現工作的分享。

這裡，我們先分享，不基於docker的，復現AlphaFold開源的推理，RoseTTAFold的推理。

1、AlphaFold2的host版本推理，Ubuntu18.04+CUDA10.1

apt-get update && apt-get -y install aria2nohup ~/alphafold/scripts/download_pdb_mmcif.sh /root/alphafold/data/pdb_mmcif/ &  wget other datasets
cd ~/
git clone https://github.com/deepmind/alphafold.gitcd alphafold

pip install dm-haiku  #jax之上最流行的NN庫pip install simtk
pip install openmmpip install absl-py==0.13.0
curl -fsSL https://mirrors.aliyun.com/nvidia-cuda/ubuntu1804/x86_64/7fa2af80.pub | apt-key add -echo "deb https://mirrors.aliyun.com/nvidia-cuda/ubuntu1804/x86_64/ /" > /etc/apt/sources.list.d/cuda.listapt updateDEBIAN_FRONTEND=noninteractive apt-get install -y build-essential cmake git hmmer kalign tzdata wget cuda-command-line-tools-10-1 rm -rf /var/lib/apt/lists/*
git clone --branch v3.3.0 http://github.com/soedinglab/hh-suite.git /tmp/hh-suitemkdir /tmp/hh-suite/build && cd /tmp/hh-suite/buildcmake -DCMAKE_INSTALL_PREFIX=/opt/hhsuite ..make -j 4 && make installln -s /opt/hhsuite/bin/* /usr/bin
wget -q -P /tmp https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.shbash /tmp/Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda

export PATH="/opt/conda/bin:$PATH"conda update -qy condaconda install -y -c conda-forge openmm=7.5.1 cudatoolkit==10.1.243 pdbfixer pip
wget -q -P ~/alphafold/alphafold/common/ https://git.scicore.unibas.ch/schwede/openstructure/-/raw/7102c63615b64735c4941278d92b554ec94415f8/modules/mol/alg/src/stereo_chemical_props.txt
pip3 install --upgrade pippip3 install -r ~/alphafold/requirements.txtpip3 install --upgrade jax jaxlib==0.1.69+cuda101 -f  https://storage.googleapis.com/jax-releases/jax_releases.html
cd /opt/conda/lib/python3.8/site-packagespatch -p0 < ~/alphafold/docker/openmm.patchcd ~/alphafoldwget https://storage.googleapis.com/alphafold/alphafold_params_2021-07-14.tar && mkdir && tar params -xvf alphafold_params_2021-07-14.tar -C params
wget -o /dev/null -O T1050.fasta 'https://www.predictioncenter.org/casp14/target.cgi?target=T1050&view=sequence'
export DATA_DIR='/root/alphafold/data/' && export OUTPUT_DIR='/root/alphafold/output/'
python run_alphafold.py --fasta_paths=${DATA_DIR}/fasta_path_0/T1050.fasta --uniref90_database_path=${DATA_DIR}/uniref90_database_path/uniref90.fasta --mgnify_database_path=${DATA_DIR}/mgnify_database_path/mgy_clusters.fa --uniclust30_database_path=${DATA_DIR}/uniclust30_database_path/uniclust30_2018_08 --bfd_database_path=${DATA_DIR}/bfd_database_path/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt --pdb70_database_path=${DATA_DIR}/pdb70_database_path/pdb70 --data_dir=${DATA_DIR}/data_dir/data --template_mmcif_dir=${DATA_DIR}/template_mmcif_dir/mmcif_files --obsolete_pdbs_path=${DATA_DIR}/obsolete_pdbs_path/obsolete.dat --output_dir=${OUTPUT_DIR}/output --model_names=model_1,model_2,model_3,model_4,model_5 --max_template_date=2020-05-14 --preset=full_dbs --benchmark=False --logtostderr

2、RoseTTAFold的推理，Ubuntu18.04+CUDA10.1
cat /usr/local/cuda/version.txtpython --versionpip install dgl-cu101
pip install torch_geometric torch_sparse torch_scatter
wget https://github.com/soedinglab/hh-suite/releases/download/v3.2.0/hhsuite-3.2.0-AVX2-Linux.tar.gz --no-check-certificatetar xvfz hhsuite-3.2.0-AVX2-Linux.tar.gzmkdir hhsuitemv bin hhsuitemv data hhsuitemv scripts hhsuitemv LICENSE hhsuitemv README hhsuite
wget ftp://ftp.ncbi.nih.gov/blast/executables/legacy.NOTSUPPORTED/2.2.26/blast-2.2.26-x64-linux.tar.gztar -xzvf blast-2.2.26-x64-linux.tar.gz
apt install gcc-5 g++-5update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-5 40update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-5 40
apt-get install libsparsehash-dev  
git clone https://github.com/soedinglab/csblastcd csblast/srcmake csblastmake csbuildmake csclustmake cssgdmake cstrainsetmake cstranslate
wget https://files.ipd.uw.edu/pub/RoseTTAFold/weights.tar.gztar xfz weights.tar.gz
git clone https://github.com/RosettaCommons/RoseTTAFoldcd RoseTTAFold
#修改run_e2e_ver.sh腳本#set -e#mkdir -p $WDIR/log#conda activate RoseTTAFold#Backer版本，有端到端版本，還有多步的版本；對步的版本，會先預測出多組候選距離與朝向，然後做精化，最終評估選擇出最好的結果修改run_pyrosetta_ver.sh腳本
#set -e#mkdir -p $WDIR/log#conda activate RoseTTAFold#conda deactivate#conda activate folding
修改input_prep/make_msa.sh#-d $MYDB    
修改input_prep/make_ss.sh中的修改相關路徑，參考如下：DATADIR="/root/data/psipred/soft/psipred/data"/root/data/psipred/soft/blast-2.2.26/bin/makemat -P $ID
cd <workspace>ln -s /root/data/hh/data/UniRef30_2020_06/ .ln -s /root/data/baker/pdb100_2021Mar03/ .
ln -s /root/data/psipred/soft/csblast/ .mv ./csblast ./csblast-2.2.3
ln -s /root/data/baker/weights/ .
執行shell腳本，開始預測，shell腳本任選一個或者都執行
#端到端預測，效果略差run_e2e_ver.sh
#pipeline多步預測run_pyrosetta_ver.sh

如上，最直白最簡單的介紹了蛋白質結構預測問題和相關比賽，對nature和science最新發表的兩篇相關前沿論文，做了簡單的技術點分析，也分享了我們在復現過程中的一些經驗。如果您對完整的復現包含訓練感興趣，可以關注我們後續的分享，包含訓練過程，分布式版本，進一步調優等工作。
MindSpore開源社區招聘進行中，詳情戳下圖了解！
MindSpore官方資料
GitHub : https://github.com/mindspore-ai/mindspore
Gitee : https : //gitee.com/mindspore/mindspore
官方QQ群 : 871543426