PSP - 蛋白质结构预测 ESMFold 算法的工程配置

欢迎关注我的CSDN：https://spike.blog.csdn.net/
本文地址：https://blog.csdn.net/caroline_wendy/article/details/129062590

ESMFold 是一种基于预训练语言模型的蛋白质结构预测方法，主要思想是利用大规模预训练蛋白质语言模型来替代 MSA（多序列比对），并且语言模型的规模与结构预测效果具有明显的正相关。ESMFold 的流程分为两个部分：序列预训练和结构预测。

• 序列预训练部分：使用了一个 15 亿参数的 Transformer 模型，叫做 ESM-2，来学习蛋白质序列自身的表征。ESM-2 采用随机 Mask 掉残基来预测被mask残基类型的训练方式，类似于 BERT。ESM-2 还对位置编码做了修改，可以支持更长的氨基酸序列编码。
• 结构预测部分：将预训练好的语言模型 ESM-2 的蛋白质序列 Embedding 和 Attention Map 接入与 48 层 Folding Trunk 和 8 层 Structure Module，来预测蛋白质全原子的结构。Folding Trunk 是退化版的 Evoformer，因为只有单序列，所以 Axis Attention 机制就退化成了普通的 Self-Attention，而节点与边 Embbeding 的更新方式保持相同。Structure Module与AlphaFold2相同。

ESMFold 在训练时不仅使用了 PDB 数据库中实验解析的单链结构，还使用了 AlphaFold2 预测的高置信度的蛋白结构。这样做可以增加训练数据量，并提高孤儿蛋白（MSA很少或没有MSA的蛋白）的结构预测精度。ESMFold 在 CASP15 和 CAMEO 上取得了和 AlphaFold2 相当的效果，并且效率大幅提升（单卡V100预测一个384长的蛋白质氨基酸序列的结构，仅耗时14.2秒，传统MSA需要半小时）。

工程：

• 论文：Language models of protein sequences at the scale of evolution enable accurate structure prediction
• GitHub: https://github.com/facebookresearch/esm

使用 FAIR 提供的服务，直接访问：

curl -X POST --data "KVFGRCELAAAMKRHGLDNYRGYSLGNWVCAAKFESNFNTQATNRNTDGSTDYGILQINSRWWCNDGRTPGSRNLCNIPCSALLSSDITASVNCAKKIVSDGNGMNAWVAWRNRCKGTDVQAWIRGCRL" https://api.esmatlas.com/foldSequence/v1/pdb/ > test.pdb

---

1. 配置环境

下载工程：

git clone [email protected]:facebookresearch/esm.git

1.1 自动安装

先直接使用 environment.yml 安装环境：

cd esm
conda env create -f environment.yml

再具体安装失败的步骤。

1.2 手动安装

准备 conda 环境，需要指定 Python 为 3.7.12，安装 ESM 相关的 Python 包，版本参考 ESMFold 的 environment.yml：

conda create -n esm python=3.7.12
conda activate esm

# 更新安装包，torch 需要低于 2.0 版本 
pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu116
pip install biopython==1.79 deepspeed==0.5.9 pytorch_lightning==1.5.10

# 安装环境包
python -m pip install --upgrade pip
# 如遇到错误，多重复几次
pip install 'dllogger @ git+https://github.com/NVIDIA/dllogger.git'
pip install 'openfold @ git+https://github.com/aqlaboratory/openfold.git@4b41059694619831a7db195b7e0988fc4ff3a307'

# 安装ESM包
pip install fair-esm
pip install git+https://github.com/facebookresearch/esm.git
pip install "fair-esm[esmfold]"

pip install biotite==0.37.0

测试 PyTorch 是否安装成功：

import torch
print(torch.__version__)  # 1.12.1
print(torch.cuda.is_available())  # True

1.3 Bugfix

Bug1: 遇到Bug “ModuleNotFoundError: No module named torch._six”

即：

File "miniconda3/envs/esm/lib/python3.9/site-packages/deepspeed/runtime/utils.py", line 18, in <module>
    from torch._six import inf
ModuleNotFoundError: No module named 'torch._six'

原因是 torch 版本与deepspeed 版本不兼容，升级 deepspeed==0.9.5 与 torch==2.0.1 版本兼容，那么 deepspeed 与 openfold版本又不兼容。

只能将 torch 降级为1.12.1，即

pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu116
pip install deepspeed==0.5.9

其他情况，也可以重新安装 DeepSpeed：

pip uninstall deepspeed
pip install deepspeed==0.8.3

Bug2: 遇到Bug “AttributeError: module ‘deepspeed.utils’ has no attribute ‘is_initialized’”

即：

File "miniconda3/envs/esm/lib/python3.9/site-packages/openfold/model/primitives.py", line 198, in forward
  deepspeed.utils.is_initialized()
AttributeError: module 'deepspeed.utils' has no attribute 'is_initialized'

当前全部版本：

ERROR: Could not find a version that satisfies the requirement deepspeed== (from versions: 0.3.1.dev1, 0.3.1.dev2, 0.3.1.dev3, 0.3.1.dev4, 0.3.1.dev5, 0.3.1.dev6, 0.3.1.dev7, 0.3.1.dev8, 0.3.1, 0.3.2, 0.3.3, 0.3.4, 0.3.5, 0.3.6, 0.3.7, 0.3.8, 0.3.9, 0.3.10, 0.3.11, 0.3.12, 0.3.13, 0.3.14, 0.3.15, 0.3.16, 0.4.0, 0.4.1, 0.4.2, 0.4.3, 0.4.4, 0.4.5, 0.5.0, 0.5.1, 0.5.2, 0.5.3, 0.5.4, 0.5.5, 0.5.6, 0.5.7, 0.5.8, 0.5.9, 0.5.10, 0.6.0, 0.6.1, 0.6.3, 0.6.4, 0.6.5, 0.6.6, 0.6.7, 0.7.0, 0.7.1, 0.7.2, 0.7.3, 0.7.4, 0.7.5, 0.7.6, 0.7.7, 0.8.0, 0.8.1, 0.8.2, 0.8.3, 0.9.0, 0.9.1, 0.9.2, 0.9.3, 0.9.4, 0.9.5)

参考 Bug1， deepspeed 与 openfold 版本不兼容，deepspeed 版本降级。

Bug3: 遇到Bug “Segmentation fault (core dumped)”

即

>>> import torch
Segmentation fault (core dumped)

原因是LD_LIBRARY_PATH路径错误，调用

unset LD_LIBRARY_PATH

即可。

Bug4: 遇到Bug “ImportError: cannot import name ‘seed_everything’ from ‘pytorch_lightning.utilities.seed’”

即 pytorch_lightning 版本错误，重新安装即可。

pip install pytorch_lightning==1.5.10

**Bug5: AttributeError: partially initialized module 'cv2' has no attribute '_registerMatType'"

参考：https://github.com/Project-MONAI/MONAILabel/discussions/869

---

2. 测试环境

测试脚本 test.py：

import torch
import esm

model = esm.pretrained.esmfold_v1()
model = model.eval().cuda()

# Optionally, uncomment to set a chunk size for axial attention. This can help reduce memory.
# Lower sizes will have lower memory requirements at the cost of increased speed.
# model.set_chunk_size(128)

sequence = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"
# Multimer prediction can be done with chains separated by ':'

with torch.no_grad():
    output = model.infer_pdb(sequence)

with open("result.pdb", "w") as f:
    f.write(output)

import biotite.structure.io as bsio
struct = bsio.load_structure("result.pdb", extra_fields=["b_factor"])
print(struct.b_factor.mean())  # this will be the pLDDT
# 88.3

运行测试脚本：

python3 test.py

输出：result.pdb，以及 pLDDT 的值是88.28875494071147，即：

参考：

• CSDN - 解决No module named ’torch._six‘问题
• GitHub - [BUG] No module named ‘torch._six’
• PyTorch - previous-versions
• OpenFold - environment.yml
• CSDN - segment fault(core dumped)