Meta AI在Science正刊发表的蛋白质结构预测模型ESMFold运行经验分享

1. 摘要

ESMFold是由Meta AI团队开发的一种基于深度学习的高效蛋白质结构预测模型，其核心目标是利用大规模蛋白质语言模型（ESM）直接从氨基酸序列快速推断蛋白质的三维结构。ESMFold通过预训练的语言模型捕捉序列中的进化与结构关联性，结合几何优化模块生成高精度原子坐标，显著降低了传统方法对多重序列比对（MSA）和模板依赖的计算成本。该模型在蛋白质从头预测（de novo prediction*）、功能位点解析、突变效应模拟等领域具有重要价值，以高效的推理性能，推动结构预测技术的普惠化应用。

2. ESMFold介绍

ESMFold是一种基于深度学习的蛋白质结构预测框架，其核心创新在于将超大规模蛋白质语言模型（如ESM-2）与几何优化模块结合，直接从氨基酸序列预测三维结构，于2023年正式发表于《Science》期刊。传统方法依赖计算密集的MSA检索与同源模板匹配，而ESMFold通过自监督学习提取序列深层语义特征，直接建模残基间距离与角度分布，大幅降低计算成本。此外，其推理能力支持大规模突变效应模拟（如酶活性位点优化）与结构组学研究。开源代码与预训练模型的发布，进一步降低了结构生物学的技术门槛，与AlphaFold2形成互补，共同拓展了计算驱动的蛋白质工程边界。

3. 网络架构

ESMFold的核心架构由两大模块构成：蛋白质语言模型编码器（ESM-2）与几何解码优化模块，其设计目标是通过端到端深度学习直接从氨基酸序列预测高精度三维结构。

图片

图1：esmfold的网络架构

3.1. 蛋白质语言模型编码器（ESM-2）

输入处理： 输入为单条氨基酸序列（长度≤1024），通过词嵌入层映射为初始特征向量（维度1280）。

注意力编码： 采用34层Transformer架构。每层包含：多头自注意力机制，用于捕捉序列内残基间的长程依赖与进化关联；残差连接与层归一化，输出为 LayerNorm(x + SelfAttention(x))；和前馈网络，其中每层前馈网络中又有两层全连接层，激活函数为GELU。

输出特征：最终提取的序列嵌入包含全局结构语义（维度2560），用于后续几何解码。

3.2. 几何解码优化模块

几何特征预测：

• 残基间距离与角度建模：通过轻量化卷积网络预测残基对的Cα-Cα距离及主链二面角。
• 侧链构象预测：基于主链几何，利用旋转平移不变网络（SE(3)-Transformer）预测侧链χ角分布。

结构生成与迭代优化：

• 初始坐标生成：通过Kabsch算法将距离矩阵转换为初始三维坐标；
• 能量最小化：结合物理力场（如范德华力、键角约束）与预测几何损失，使用梯度下降迭代优化结构。

3.3. 关键技术优化

计算效率提升：

• 舍弃传统MSA检索，仅依赖单序列输入；
• 几何解码模块采用低秩近似与稀疏注意力，单GPU推理时间缩短至数分钟（如400残基蛋白约2分钟）。

训练策略：

• 两阶段训练：先预训练ESM-2语言模型（250M参数，UniRef50数据集），再联合微调解码模块
• 损失函数：联合优化距离分类交叉熵、角度均方误差（MSE）及结构相似性指标（如TM-score）。

4. 实验

4.1. 起容器

docker run  -it \
--privileged=true \
--device /dev/davinci0 \
--device /dev/davinci1 \
--device /dev/davinci2 \
--device /dev/davinci3 \
--device /dev/davinci4 \
--device /dev/davinci5 \
--device /dev/davinci6 \
--device /dev/davinci7 \
--device  /dev/davinci_manager \
--device  /dev/devmm_svm \
--device /dev/hisi_hdc \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/bin/hccn_tool:/usr/bin/hccn_tool \
-v /usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/driver/lib64/common \
-v /usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64/driver \
-v /etc/ascend_install.info:/etc/ascend_install.info \
-v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
--name esmfold 8f1cb37dd88c  /bin/bash

其镜像内已经预安装了torch的相关组件

4.2. 安装esmfold

pip install fair-esm

图片

pip install git+https://github.com/facebookresearch/esm.git

图片

注：这两种方式都是安装esmfold的，选其中一个执行即可

报错没有git，需要安装git工具

apt-get install -y git

图片

安装前需要先更新下载工具apt-get

apt-get update

图片

更新完成

apt-get install -y git

图片

下载完毕

查看git是否安装成功

git –version

图片

重新使用第二种方式安装esmfold

pip install git+https://github.com/facebookresearch/esm.git

图片

pip install "fair-esm[esmfold]"

图片

图片

图片

4.3. 安装openfold

参考文章 openfold

4.4. 验证

4.4.1. 新建测试用例

touch test.fasta
vi test.fasta

将下述蛋白质序列粘贴至test.fasta中

>test
MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG

4.4.2. 测试&处理报错

esm-fold -i test.fasta -o ./

图片

报错网络连接超时，需要手动下载权重文件

Downloading: "https://dl.fbaipublicfiles.com/fair-esm/models/esmfold_3B_v1.pt" to /root/.cache/torch/hub/checkpoints/esmfold_3B_v1.pt
Downloading: "https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t36_3B_UR50D.pt" to /root/.cache/torch/hub/checkpoints/esm2_t36_3B_UR50D.pt
Downloading: "https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t36_3B_UR50D-contact-regression.pt" to /root/.cache/torch/hub/checkpoints/esm2_t36_3B_UR50D-contact-regression.pt

需要将下载好的权重手动复制到/root/.cache/torch/hub/checkpoints/内，例：

cp esm2_t36_3B_UR50D-contact-regression.pt /root/.cache/torch/hub/checkpoints/

再次执行

esm-fold -i test.fasta -o ./

图片

报错默认使用的cuda

vi /usr/local/python3.9.2/lib/python3.9/site-packages/esm/scripts/fold.py 

导入torch_npu

from torch_npu.contrib 
import transfer_to_npu

会自动将cuda api替换为torch_npu的api

图片

esm-fold -i test.fasta -o ./

图片

图片

此时虽然还有报错，但已经能看到esfold的npu显存占用了

图片

报错AttributeError: module 'deepspeed.utils' has no attribute 'is_initialized' , 原因是由于 openfold 代码调用了旧版 deepspeed 的 api，而新版 deepspeed 中该api 已被移除。

需要降低deepspeed的版本

pip install deepspeed==0.9.5

图片

报错deepspeed 0.9.5 与新版 Pydantic（V2+）不兼容

deepspeed 0.9.5 依赖旧版 Pydantic（V1），当前环境中安装的 Pydantic 是 V2+ 版本，导致 FieldInfo 对象缺少 required 属性，所以需要降低Pydantic版本。

pip install pydantic==1.10.13
pip install deepspeed==0.9.5

图片

esm-fold -i test.fasta -o ./

还是报openfold和deepspeed版本冲突的错

图片

pip install deepspeed==0.10.3版本

这个版本明确知道有is_initialized这个属性

图片

esm-fold -i test.fasta -o ./

依旧报错

图片

图片

怀疑is_initialized属性不在deepspeed.utils模块

grep -rn is_initialized查看在哪定义的

图片

果然不在utils内定义，是在comm模块定义的

验证一下

图片

修改源码

vi /usr/local/python3.9.2/lib/python3.9/site-packages/openfold/model/primitives.py

将199行的deepspeed.utils.is_initialized()修改为deepspeed.comm.is_initialized()

图片

保存退出

esm-fold -i test.fasta -o ./

图片

一样的问题

vi /usr/local/python3.9.2/lib/python3.9/site-packages/openfold/model/primitives.py

图片

4.4.3. 成功推理

再次执行

esm-fold -i test.fasta -o ./

图片

5. 参考文献

Lin, Z., Akin, H., Rao, R. et al. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science 380, eadf2574 (2023).

DOI: 10.1126/science.ade2574