基于昇腾适配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已不使用此写法

解决办法：

在代码中搜索所有np.object和np.bool，将其替换为object、bool
grep -rn "np\.object" openfold/
grep -rn "np\.bool" openfold/
1.
2.

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

执行推理

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

重新执行推理

推理成功

输出文件

-----------------------------------