基于昇腾适配 DeepMind 团队发布的蛋白质结构预测模型 OpenFold

作者：Splendid2025

2025-06-13
浙江
本文字数：6377 字
阅读完需：约 21 分钟

1. 摘要

OpenFold 是一种基于深度学习的蛋白质结构预测模型，广泛应用于蛋白质从头预测、功能位点解析、突变效应模拟等领域。该模型的核心目标是通过大规模预训练和多阶段优化，从氨基酸序列中高效、准确地推断蛋白质的三维结构。OpenFold 结合了 Transformer 架构和几何优化模块，显著提高了结构预测的精度和速度。该模型的部署包含详细的微调教程、模型训练、推理优化等内容，为研究人员提供了全面的技术支持。

2. OpenFold 介绍

OpenFold 是由 DeepMind 团队开发的一种高效蛋白质结构预测模型。该模型在 AlphaFold2 的基础上进行了多项改进，进一步提升了蛋白质结构预测的准确性和计算效率。其核心算法包括大规模预训练的 Transformer 模型和几何优化模块，能够从氨基酸序列中快速推断出蛋白质的三维结构。通过多阶段优化和大规模数据集的训练，该模型在蛋白质从头预测、功能位点解析、突变效应模拟等领域展现了卓越的性能。此外，OpenFold 的部署文档详细介绍了模型的微调、训练、推理优化等步骤，为研究人员提供了全面的技术支持，推动了蛋白质结构预测技术的广泛应用。

3. OpenFold 网络架构

OpenFold 的模型架构由三个核心模块构成：输入嵌入层、Evoformer 堆叠模块和结构解码器。输入数据整合了多序列比对（MSA）、模板特征、氨基酸序列及进化信息，形成高维生物特征张量。通过分阶段嵌入与特征融合，数据首先被压缩至低维隐空间，随后由多尺度 Evoformer 模块进行全局-局部特征交互，最终通过几何约束的结构解码器输出蛋白质的 3D 原子坐标与置信度。

3.1 输入嵌入层

为统一处理异构生物特征并降低计算复杂度，OpenFold 采用混合嵌入策略：

MSA 嵌入：使用 1D 卷积核（宽度=3，步长=1）对 MSA 序列进行通道压缩，配合层归一化（LayerNorm）稳定训练。
模板嵌入：通过残差连接的 3D 卷积（核 3×3×3，步长 1×2×2）提取模板结构特征，输出通道数对齐主嵌入空间。
序列特征投影：氨基酸物理化学属性经全连接层映射至隐空间，与上述嵌入结果拼接，形成初始隐状态张量（维度：C×L，L 为序列长度）。

3.2 Evoformer 堆叠模块

该模块由 48 层对称 Evoformer 块构成，采用双路处理机制：

全局注意力通路：引入轴向注意力机制，在序列维度（L）和 MSA 行维度（N）交替执行缩放余弦注意力，计算效率较传统 Transformer 提升 3.2 倍。每层包含：
局部结构通路：使用门控卷积网络（核大小=5，膨胀率=2）捕获局部氨基酸环境特征，配合三角更新机制建模残基间几何关系。每层输出经 GroupNorm 归一化后与全局通路特征融合。

3.3 结构解码器

主干几何生成：基于隐变量通过迭代对齐层（Invariant Point Attention, IPA）逐步优化主链扭转角
侧链重建：采用条件随机场（CRF）对侧链构象进行能量最小化采样，结合 Rosetta 能量函数约束立体化学合理性。
输出层：最终通过 SE(3)-等变全连接层输出原子坐标（维度：L×37×3，37 为每个残基原子数）及置信度热图（分辨率 1Å）。

4. 核心组件安装

4.1 组件版本

hdk：24.1.0.3cann：8.0.RC3python：3.9.2torch：2.1.0torch_npu：2.1.0.post6openfold：1.0.0torchaudio：2.1.0torchmetrics：1.7.1torchvision：0.16.0pytorch-lightning：1.6.5

复制代码

4.2 起容器

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  openfold  27913b525135  /bin/bash

复制代码

4.3 安装 Openfold

注：如果直接 git clone 安装，则默认安装的是 2.0.0 版本的 openfold，我们需要的是 1.0.0 版本

4.3.1 下载源码

git clone --filter=blob:none --quiet https://github.com/aqlaboratory/openfold.git ./openfold

复制代码

4.3.2 将版本修改为 1.0.0

cd ./openfold/git rev-parse -q --verify 'sha^4b41059694619831a7db195b7e0988fc4ff3a307'

复制代码

git fetch -q https://github.com/aqlaboratory/openfold.git 4b41059694619831a7db195b7e0988fc4ff3a307

复制代码

git checkout -q 4b41059694619831a7db195b7e0988fc4ff3a307

复制代码

du -sh

复制代码

再 vi setup.py 查看，此时就变成了 1.0.0 版本

4.3.3 修改 openfold 的 setup.py 文件

vi /home/openfold/setup.py

复制代码

在头部从 torch.utils.cpp_extension 中增加对 cppextension 的引用

修改 get_cuda_bare_metal_version 函数，增加对有没有 cuda 的判断

修改 ext_modules 的内容如下

4.3.4 编译安装

pip install e .

复制代码

报错找不到 torch，但此时 torch 已经安装好了

pip install --upgrade setuptools

复制代码

更新 setuptools 工具

pip install --upgrade pippip install wheel

复制代码

更新 pip，安装 wheel

pip install e .

复制代码

报错

但 pip list 却能查到 openfold（可能是 python 自动将当前目录加入了模块搜索路径导致的）

看能否导入

python -c "import openfold; print(openfold.__version__)"

复制代码

报错 ModuleNotFoundError: No module named 'torch._six'

原因：pytorch 版本与 deepspeed 版本冲突，torch._six 是 pytorch 旧版本中的模块，新版本已移除。

解决办法：更新 deepspeed

pip install --upgrade deepspeed

复制代码

报错

设置环境变量

export ASCEND_HOME_PATH=/usr/local/Ascend/ascend-toolkit/pip install --upgrade deepspeed

复制代码

继续报错

source cannsource /usr/local/Ascend/ascend-toolkit/set_env.shpip install --upgrade deepspeed

复制代码

python -c "import openfold; print(openfold.__version__)"

复制代码

报错原因：pytorch lightning 版本不兼容，seed_everything 在当前 pytorch lightning 版本中已被移动到新的模块路径。

解决办法：降低 pytorch lightning 版本

pip install pytorch_lightning==1.6.5

复制代码

但要先降低 pip 的版本

pip install pip==23.1

复制代码

pip install pytorch_lightning==1.6.5

复制代码

python -c "import openfold; print(openfold.__version__)"

复制代码

报错原因：dllogger 是 cuda 的日志记录工具，安装需要基于 cuda 驱动，故在昇腾上面无法安装

解决办法：注销 dllogger 的导入，用 python 的标准库 logging 取代

vi /home/910_tools/openfold/openfold/utils/logger.py

复制代码

python -c "import openfold; print(openfold.__version__)"

复制代码

报错 ModuleNotFoundError: No module named 'attn_core_inplace_cuda'

pip list | grep open

复制代码

此时只有在 openfold 的源码目录才能用 pip list 查询到

且无法卸载

怀疑根本就没装好，重新安装

export OPENFOLD_DISABLE_CUDA=1 # 运行时禁用CUDA相关代码export NO_CUDA=1 # 安装时禁用CUDA扩展编译pip install . --no-build-isolation

复制代码

报错 ninja: error: '/home /910_tools/openfold/openfold/utils/kernel/csrc/softmax_cuda_stub.cpp', needed by '/home /910_tools/openfold/build/temp.linux-aarch64-cpython-39/openfold/utils/kernel/csrc/softmax_cuda_stub.o', missing and no known rule to make it

vi /home/910_tools/openfold/setup.py

复制代码

将 cmdclass={'build_ext': BuildExtension}改为 cmdclass={'build_ext': BuildExtension.with_options(use_ninja=False)}

pip install . --no-build-isolation

复制代码

报错 g++: error: openfold/utils/kernel/csrc/softmax_cuda_stub.cpp: No such file or directory

查看 csrc 文件夹

ll /home /910_tools/openfold/openfold/utils/kernel/csrc

复制代码

发现确实没有 softmax_cuda_stub.cpp 这个文件

但 GitHub 上面的源码却有这个文件

touch softmax_cuda_stub.cpp

vi softmax_cuda_stub.cpp

将源码复制进去 wq 保存退出

再次执行

pip install . --no-build-isolation

复制代码

安装成功

验证

此时不论哪里都能查得到 openfold

4.4 安装 OpenMM

直接 pip install openmm 会报错

报错原因：当前服务器架构是 aarch64，而通过 pypi 下载的 openmm，没有 aarch64 架构的 whl 离线包，所以会安装失败。

pypi.org 上面仅有 8.1.1 版本以上且没有 aarch64 架构的包

通过源码编译的方式安装失败

解决办法

在 conda-forge 的 package 中搜索 openmm

选择要安装的版本和机器架构

下载 openmm 的预编译的二进制包（非源码包）

注: 如果是 conda 环境，直接在当前环境执行

conda install --use-local openmm-7.7.0-py39h127581e_1.tar.bz2

复制代码

若是 docker 环境还需以下操作

解压

tar -xjf openmm-7.7.0-py39h127581e_1.tar.bz2

复制代码

将 openmm 的 python 模块复制到容器中 python 的 site-packages 中

cp -r ./lib/python3.9/site-packages/openmm   /usr/local/python3.9.2/lib/python3.9/site-packages/

复制代码

将 openmm 的共享库（.so 文件）复制到 Python 的 lib 目录并更新动态链接器缓存

cp libOpenMM*  /usr/local/python3.9.2/lib/ldconfig

复制代码

验证容器的 python 环境中 openmm 的安装情况

4.5 安装 pdbfixer

直接 pip install pdbfixer 会报错

pip install pdbfixer

复制代码

下载源码使用 pip install .安装

git clone https://github.com/openmm/pdbfixer.git

复制代码

pip install .

复制代码

但此时默认安装的是最新版的 pdbfixer，提示需要 8.2 版本以上的 openmm

查看所有版本信息

git ls-remote --tags origin

复制代码

从源码中看到 1.8.1 版本的 pdbfixer，要求 openmm 的版本大于 7.2 即可，所以安装 1.8.1 版本

切换至 1.8.1 版本

git checkout tags/v1.8.1

复制代码

pip install .

复制代码

安装成功

pip list | grep pdb

复制代码

查看

5. 实验

5.1 下载 openfold 权重

bash scripts/download_openfold_params.sh openfold/resources

复制代码

报错，需要安装 aws

安装好之后重新执行

报错

原因：在 Python 3 环境中运行了为 Python 2 编写的代码。print 语句缺少括号。

解决办法：卸载 aws，重新安装。

curl "https://awscli.amazonaws.com/awscli-exe-linux-aarch64.zip" -o "awscliv2.zip"

复制代码

./aws/installaws –version

复制代码

重新执行脚本下载权重

bash scripts/download_openfold_params.sh openfold/resources

复制代码

5.2 下载数据集

bash scripts/download_alphafold_dbs.sh data/

复制代码

报错 Error: aria2c could not be found. Please install aria2c (sudo apt install aria2).

apt update && apt install aria2 -y

复制代码

开始下载

但原始数据集太大，手动创建蛋白质序列测试

5.3 创建蛋白质序列测试文件

新建 data 目录，在 data 目录新建 test.fasta

vi test.fasta

复制代码

内容为

>test_sequence|PDBID|1AKI

GIVEQCCTSICSLYQLENYCN

保存退出

注：上述为简化胰岛素类似物（PDB ID: 1AKI）蛋白质序列

5.4 下载蛋白质模板文件

mkdir -p template_mmcif_dir/pdb_mmcif/mmcif_files/

复制代码

下载蛋白质模板文件

https://files.wwpdb.org/pub/pdb/data/structures/divided/mmCIF/ak/1aki.cif.gz

复制代码

解压之后放至 template_mmcif_dir/pdb_mmcif/mmcif_files/

5.5 推理

vi run_pretrained_openfold.py

复制代码

导入 torch-npu

执行推理

python3 run_pretrained_openfold.py  /home/910_tools/data/  /home/910_tools/template_mmcif_dir/pdb_mmcif/mmcif_files/  --output_dir /home/910_tools/output/ --model_device "npu:0" --config_preset "model_1_ptm"  --openfold_checkpoint_path /home/openfold/openfold/resources/openfold_params/finetuning_ptm_2.pt

复制代码

报错

原因：openfold 中使用了 np.object 的方式，而新版本 numpy 已不使用此写法

解决办法：

在代码中搜索所有 np.object 和 np.bool，将其替换为 object、bool

grep -rn "np\.object" openfold/grep -rn "np\.bool" openfold/

复制代码

2. 降低 numpy 版本，但要考虑更多的版本冲突（不建议）

执行推理

python3 run_pretrained_openfold.py  /home/910_tools/data/  /home/910_tools/template_mmcif_dir/pdb_mmcif/mmcif_files/  --output_dir /home/910_tools/output/ --model_device "npu:0" --config_preset "model_1_ptm"  --openfold_checkpoint_path /home/openfold/openfold/resources/openfold_params/finetuning_ptm_2.pt

复制代码

报错

报错原因：deepspeed 属性不在 utils 模块

解决办法：将涉及到 deepspeed.utils.is_initialized()的地方都修改为 deepspeed.comm.is_initialized()

执行推理

报错

报错原因：openfold 自定义了一个注意力机制，底层会调用 cuda 核心进行加速

解决办法：使用原生的 torch 操作代替 cuda 操作

vi /home/910_tools/openfold/openfold/utils/kernel/attention_core.py

复制代码

删除原始代码

将下面代码复制到 attention_core.py 保存退出

import torchimport torch.nn.functional as Ffrom functools import reducefrom operator import mulclass AttentionCoreFunction(torch.autograd.Function):    @staticmethod    def forward(ctx, q, k, v, bias_1=None, bias_2=None):        if bias_1 is None and bias_2 is not None:            raise ValueError("bias_1 must be specified before bias_2")        q = q.contiguous()        k = k.contiguous()        v = v.contiguous()        attention_logits = torch.matmul(q, k.transpose(-1, -2))        if bias_1 is not None:            attention_logits += bias_1        if bias_2 is not None:            attention_logits += bias_2        attention_probs = F.softmax(attention_logits, dim=-1)        o = torch.matmul(attention_probs, v)        ctx.save_for_backward(q, k, v, attention_probs)        ctx.bias_1_shape = bias_1.shape if bias_1 is not None else None        ctx.bias_2_shape = bias_2.shape if bias_2 is not None else None        return o    @staticmethod    def backward(ctx, grad_output):        q, k, v, attention_probs = ctx.saved_tensors        grad_q = grad_k = grad_v = grad_bias_1 = grad_bias_2 = None        grad_v = torch.matmul(attention_probs.transpose(-1, -2), grad_output)        grad_attn = torch.matmul(grad_output, v.transpose(-1, -2))        grad_attn_logits = attention_probs * (grad_attn - (grad_attn * attention_probs).sum(dim=-1, keepdim=True))        grad_q = torch.matmul(grad_attn_logits, k)        grad_k = torch.matmul(q.transpose(-1, -2), grad_attn_logits).transpose(-1, -2)        if ctx.bias_1_shape is not None:            grad_bias_1 = grad_attn_logits.sum(                dim=tuple(i for i, d in enumerate(ctx.bias_1_shape) if d == 1),                keepdim=True            )        if ctx.bias_2_shape is not None:            grad_bias_2 = grad_attn_logits.sum(                dim=tuple(i for i, d in enumerate(ctx.bias_2_shape) if d == 1),                keepdim=True            )        return grad_q, grad_k, grad_v, grad_bias_1, grad_bias_2attention_core = AttentionCoreFunction.apply

复制代码

执行推理

报错，和上述报错一本质一样

vi /home/910_tools/openfold/openfold/model/structure_module.py

复制代码

注销 45 行 attn_core_inplace_cuda = importlib.import_module("attn_core_inplace_cuda")

添加 from openfold.utils.kernel.attention_core import AttentionCoreFunction

删除 435 行调用 attn_core_inplace_cuda.forward_代码，使用原生 softmax 实现

执行推理

报错 ModuleNotFoundError: No module named 'simtk.openmm'

原因：openmm 或相关依赖中使用比较老的 openmm 导入方式

解决办法：将 openfold、openmm、pdbfixer 中所有涉及到 simtk.openmm 的地方都修改为 openmm

重新执行推理

推理成功

输出文件

发布于: 刚刚阅读数: 4

Splendid2025

关注

还未添加个人签名 2025-01-26 加入

AI4SCI

发布

暂无评论

创作场景

基于昇腾适配 DeepMind 团队发布的蛋白质结构预测模型 OpenFold

1. 摘要

2. OpenFold 介绍

3. OpenFold 网络架构

3.1 输入嵌入层

3.2 Evoformer 堆叠模块

3.3 结构解码器

4. 核心组件安装

4.1 组件版本

4.2 起容器

4.3 安装 Openfold

4.4 安装 OpenMM

4.5 安装 pdbfixer

5. 实验

5.1 下载 openfold 权重

5.2 下载数据集

5.3 创建蛋白质序列测试文件

5.4 下载蛋白质模板文件

5.5 推理

Splendid2025

评论