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

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.3
cann：8.0.RC3
python：3.9.2
torch：2.1.0
torch_npu：2.1.0.post6
openfold：1.0.0
torchaudio：2.1.0
torchmetrics：1.7.1
torchvision：0.16.0
pytorch-lightning：1.6.5

• 1.
• 2.
• 3.
• 4.
• 5.
• 6.
• 7.
• 8.
• 9.
• 10.

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

• 1.
• 2.
• 3.
• 4.
• 5.
• 6.
• 7.
• 8.
• 9.
• 10.
• 11.
• 12.
• 13.
• 14.
• 15.
• 16.
• 17.
• 18.
• 19.
• 20.
• 21.

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

• 1.

4.3.2 将版本修改为1.0.0

登录后复制

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

• 1.
• 2.

登录后复制

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

• 1.

登录后复制

git checkout -q 4b41059694619831a7db195b7e0988fc4ff3a307

• 1.

登录后复制

du -sh

• 1.

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

4.3.3 修改openfold的setup.py文件

登录后复制

vi /home/openfold/setup.py

• 1.

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

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

修改ext_modules的内容如下

4.3.4 编译安装

登录后复制

pip install e .

• 1.

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

登录后复制

pip install --upgrade setuptools

• 1.

更新setuptools工具

登录后复制

pip install --upgrade pip
pip install wheel

• 1.
• 2.

更新 pip，安装wheel

登录后复制

pip install e .

• 1.

报错

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

看能否导入

登录后复制

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

• 1.

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

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

解决办法：更新deepspeed

登录后复制

pip install --upgrade deepspeed

• 1.

报错

设置环境变量

登录后复制

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

• 1.
• 2.

继续报错

登录后复制

source cann
source /usr/local/Ascend/ascend-toolkit/set_env.sh
pip install --upgrade deepspeed

• 1.
• 2.
• 3.

登录后复制

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

• 1.

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

解决办法：降低pytorch lightning版本

登录后复制

pip install pytorch_lightning==1.6.5

• 1.

但要先降低pip的版本

登录后复制

pip install pip==23.1

• 1.

登录后复制

pip install pytorch_lightning==1.6.5

• 1.

登录后复制

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

• 1.

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

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

登录后复制

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

• 1.

登录后复制

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

• 1.

报错ModuleNotFoundError: No module named 'attn_core_inplace_cuda'

登录后复制

pip list | grep open

• 1.

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

且无法卸载

在根目录或者其它目录查询不到

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

登录后复制

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

• 1.
• 2.
• 3.

报错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

• 1.

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

登录后复制

pip install . --no-build-isolation

• 1.

报错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

• 1.

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

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

touch softmax_cuda_stub.cpp

vi softmax_cuda_stub.cpp

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

再次执行

登录后复制

pip install . --no-build-isolation

• 1.

安装成功

验证

此时不论哪里都能查得到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

• 1.

若是docker环境还需以下操作

解压

登录后复制

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

• 1.

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

登录后复制

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

• 1.

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

登录后复制

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

• 1.
• 2.

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

4.5 安装pdbfixer

直接pip install pdbfixer会报错

登录后复制

pip install pdbfixer

• 1.

下载源码使用pip install .安装

登录后复制

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

• 1.

登录后复制

pip install .

• 1.

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

查看所有版本信息

登录后复制

git ls-remote --tags origin

• 1.

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

切换至1.8.1版本

登录后复制

git checkout tags/v1.8.1

• 1.

登录后复制

pip install .

• 1.

安装成功

登录后复制

pip list | grep pdb

• 1.

查看

5. 实验

5.1 下载openfold权重

登录后复制

bash scripts/download_openfold_params.sh openfold/resources

• 1.

报错，需要安装aws

安装好之后重新执行

报错

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

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

登录后复制

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

• 1.

登录后复制

./aws/install
aws –version

• 1.
• 2.

重新执行脚本下载权重

登录后复制

bash scripts/download_openfold_params.sh openfold/resources

• 1.

5.2 下载数据集

登录后复制

bash scripts/download_alphafold_dbs.sh data/

• 1.

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

登录后复制

apt update && apt install aria2 -y

• 1.

开始下载

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

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

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

登录后复制

vi test.fasta

• 1.

内容为

>test_sequence|PDBID|1AKI

GIVEQCCTSICSLYQLENYCN

保存退出

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

5.4 下载蛋白质模板文件

登录后复制

mkdir -p template_mmcif_dir/pdb_mmcif/mmcif_files/

• 1.

下载蛋白质模板文件

登录后复制

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

• 1.

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

5.5 推理

登录后复制

vi run_pretrained_openfold.py

• 1.

导入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

• 1.

报错

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

解决办法：

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

登录后复制

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

• 1.
• 2.

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

• 1.

报错

报错原因：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

• 1.

删除原始代码

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

登录后复制

import torch
import torch.nn.functional as F
from functools import reduce
from operator import mul
class 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_2
attention_core = AttentionCoreFunction.apply

• 1.
• 2.
• 3.
• 4.
• 5.
• 6.
• 7.
• 8.
• 9.
• 10.
• 11.
• 12.
• 13.
• 14.
• 15.
• 16.
• 17.
• 18.
• 19.
• 20.
• 21.
• 22.
• 23.
• 24.
• 25.
• 26.
• 27.
• 28.
• 29.
• 30.
• 31.
• 32.
• 33.
• 34.
• 35.
• 36.
• 37.
• 38.
• 39.
• 40.
• 41.
• 42.
• 43.
• 44.

执行推理

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

登录后复制

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

• 1.

注销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

重新执行推理

推理成功

输出文件