基于昇腾适配 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.fastavi test.fasta

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

>testMKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG

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.ptDownloading: "https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t36_3B_UR50D.pt" to /root/.cache/torch/hub/checkpoints/esm2_t36_3B_UR50D.ptDownloading: "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.13pip 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