alphafold2_jax预测蛋白质三维结构的深度学习模型
AF2
论文
模型结构
模型核心是一个基于Transformer架构的神经网络，包括两个主要组件：Sequence to Sequence Model和Structure Model，这两个组件通过迭代训练进行优化，以提高其预测准确性。
算法原理
AlphaFold2通过从蛋白质序列和结构数据中提取信息，使用神经网络模型来预测蛋白质三维结构。
环境配置
提供光源拉取推理的docker镜像：
docker pull image.sourcefind.cn:5000/dcu/admin/base/custom:alphafold2-2.3.2-dtk23.10-py38
# <Image ID>用上面拉取docker镜像的ID替换
# <Host Path>主机端路径
# <Container Path>容器映射路径
docker run -it --name alphafold --privileged --shm-size=32G --device=/dev/kfd --device=/dev/dri/ --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --ulimit memlock=-1:-1 --ipc=host --network host --group-add video -v <Host Path>:<Container Path> <Image ID> /bin/bash
镜像版本依赖：
- DTK驱动：dtk23.10
- Jax: 0.3.25
- TensorFlow2: 2.11.0
- python: python3.8
数据集
推荐使用AlphaFold2中的开源数据集，包括BFD、MGnify、PDB70、Uniclust、Uniref90等,数据集大小约2.62TB。数据集格式如下：
$DOWNLOAD_DIR/
bfd/
bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt_hhm.ffindex
bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt_hhm.ffdata
bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt_cs219.ffindex
...
mgnify/
mgy_clusters_2022_05.fa
params/
params_model_1.npz
params_model_2.npz
params_model_3.npz
...
pdb70/
pdb_filter.dat
pdb70_hhm.ffindex
pdb70_hhm.ffdata
...
pdb_mmcif/
mmcif_files/
100d.cif
101d.cif
101m.cif
...
obsolete.dat
pdb_seqres/
pdb_seqres.txt
small_bfd/
bfd-first_non_consensus_sequences.fasta
uniref30/
UniRef30_2021_03_hhm.ffindex
UniRef30_2021_03_hhm.ffdata
UniRef30_2021_03_cs219.ffindex
...
uniprot/
uniprot.fasta
uniref90/
uniref90.fasta
此处提供了一个脚本download_all_data.sh用于下载使用的数据集和模型文件：
./scripts/download_all_data.sh 数据集下载目录
数据集快速下载中心：SCNet AIDatasets ，项目中数据集可从快速下载通道下载：alphafold 。
推理
分别提供了基于Jax的单体和多体的推理脚本.
# 进入工程目录
cd alphafold2_jax
单体
./run_monomer.sh
单体推理参数说明：download_dir为数据集下载目录，monomer.fasta为推理的单体序列；--output_dir为输出目录；model_names为推理的模型名称，--model_preset=monomer为单体模型配置；--run_relax=true为进行relax操作；--use_gpu_relax=true为使用gpu进行relax操作（速度更快，但可能不太稳定），--use_gpu_relax=false为使用CPU进行relax操作（速度慢，但稳定）；若添加--use_precomputed_msas=true则可以加载已有的MSAs，否则默认运行MSA工具。
多体
./run_multimer.sh
多体推理参数说明：multimer.fasta为推理的多体序列，--model_preset=multimer为多体模型配置；--num_multimer_predictions_per_model为每个模型预测数量，其他参数同单体推理参数说明一致。
result
--output_dir目录结构如下：
<target_name>/
features.pkl
ranked_{0,1,2,3,4}.pdb
ranking_debug.json
relaxed_model_{1,2,3,4,5}.pdb
result_model_{1,2,3,4,5}.pkl
timings.json
unrelaxed_model_{1,2,3,4,5}.pdb
msas/
bfd_uniclust_hits.a3m
mgnify_hits.sto
uniref90_hits.sto
...
红色为真实结构，蓝色为预测结构
精度
测试数据：casp15、uniprot， 使用的加速卡:1张 Z100L-32G
1、plddts/iptm+ptm
单体见<target_name>/ranking_debug.json中的plddts，多体见<target_name>/ranking_debug.json中的iptm+ptm
2、其它精度值计算：TM-score: Quantitative assessment of similarity between protein structures
准确性数据：
| 数据类型 | 序列类型 | 序列 | 长度 | GDT-TS | GDT-HA | plddts/iptm+ptm | TM score | MaxSub | RMSD |
|---|---|---|---|---|---|---|---|---|---|
| fp32 | 单体 | T1029 | 125 | 0.434 | 0.256 | 93.984 | 0.471 | 0.297 | 7.202 |
| fp32 | 单体 | T1024 | 408 | 0.664 | 0.470 | 87.076 | 0.829 | 0.518 | 3.516 |
| fp32 | 多体 | H1106 | 236 | 0.203 | 0.144 | 0.860 | 0.181 | 0.151 | 20.457 |
应用场景
算法类别
蛋白质预测
热点应用行业
医疗,科研,教育
预训练权重
预训练权重快速下载中心：SCNet AIModels ，项目中的预训练权重可从快速下载通道下载：alphafold 。