OpenFold2.0 基于NPU的推理适配与测试

摘要

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

OpenFold介绍

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

OpenFold网络架构

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

环境搭建

#下载安装包
wget https://repo.anaconda.com/archive/Anaconda3-2024.10-1-Linux-aarch64.sh
chmod +x Anaconda3-2024.10-1-Linux-aarch64.sh #安装
bash Anaconda3-2024.10-1-Linux-aarch64.sh #设置环境变量
vim ~/.bashrc
export PATH=/home/openfold/anaconda3/bin:$PATH :wq
source ~/.bashrc
#创建 python3.9 环境，并激活环境
conda init
conda create -n openfold_env python=3.9
conda activate openfold_env
#配置 conda 的 lib 库
vim ~/.bashrc
export LD_LIBRARY_PATH=/home/openfold/anaconda3/lib:$LD_LIBRARY_PATH :wq
source ~/.bashrc

下载代码及安装依赖 

#下载代码
cd /home/openfold
git clone https://gitee.com/ascend/ModelZoo-PyTorch.git
#安装依赖
cd /home/openfold/ModelZoo-PyTorch/PyTorch/built-in/others/OpenFold_for_PyTorch
pip3 install -r 2.1_requirements.txt -i https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
pip install git+https://github.com/NVIDIA/dllogger.git
pip install torch==2.1.0 -i https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple #下载安装 torch_npu
wget https://gitee.com/ascend/pytorch/releases/download/v6.0.0.1-pytorch2.1.0/torch_npu- 2.1.0.post11-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
pip3 install torch_npu-2.1.0.post11-cp39-cp39- manylinux_2_17_aarch64.manylinux2014_aarch64.whl
#安装其它依赖
conda install -c bioconda -c conda-forge hhsuite
conda install -c bioconda kalign2
conda install -c bioconda -c conda-forge mmseqs2
conda install -c conda-forge openmm=7.7.0
conda install -c conda-forge pdbfixer
#构建安装 openfold
bash scripts/install_third_party_dependencies.sh python setup.py install

数据集准备

#安装 aws
curl "https://awscli.amazonaws.com/awscli-exe-linux-aarch64.zip" -o "awscliv2.zip"
unzip awscliv2.zip sudo ./aws/install #获取 pdb_mmcif 数据集 mkdir pdb_data
aws s3 cp s3://openfold/pdb_mmcif.zip pdb_data/ --no-sign-request #获取 duplicate_pdb_chains.txt
aws s3 cp s3://openfold/duplicate_pdb_chains.txt . --no-sign-request 解压 pdb_data
cd pdb_data
unzip pdb_mmcif.zip -d pdb_data
#下载参数
bash scripts/download_alphafold_params.sh openfold/resources
#修改 param 路径
vim /home/openfold/ModelZoo-PyTorch/PyTorch/built- in/others/OpenFold_for_PyTorch/examples/monomer/../../run_pretrained_openfold.py
line 490 改为:
"/home/openfold/ModelZoo-PyTorch/PyTorch/built-in/others/OpenFold_for_PyTorch/openfold", "resources", "params",
#修改 inference.sh
#修改数据集路径和 device
cd examples/monome
vim inference.sh
export MMCIF_DIR=../../pdb_data/mmcif_files/ --model_device "npu:0"

推理测试

#执行推理脚本 
bash inference.sh

报错与解决方法

1、报错 ImportError: /usr/lib64/libstdc++.so.6: version `GLIBCXX_3.4.29' not found (required by /home/openfold/anaconda3/envs/openfold_env/lib/python3.9/site- packages/numpy/fft/_pocketfft_umath.cpython-39-aarch64-linux-gnu.so) 

#更新 numpy 版本
pip3 uninstall numpy
pip3 install numpy==1.24.0

2、报错 ImportError: /usr/lib64/libstdc++.so.6: version `GLIBCXX_3.4.29' not found (required by /home/openfold/anaconda3/envs/openfold_env/lib/python3.9/site- packages/scipy/spatial/_ckdtree.cpython-39-aarch64-linux-gnu.so) 

pip3 uninstall scipy
pip install --no-binary :all: scipy

3、报错 aria2c: /home/openfold/anaconda3/lib/libstdc++.so.6: version `GLIBCXX_3.4.30' not found (required by /home/openfold/anaconda3/envs/openfold_env/bin/../lib/././libicuuc.so.75) 

conda install -c conda-forge libstdcxx-ng=12.2.0

4、报错:There is no registered Platform called “CUDA” ValueError:Minimization failed after 100 attempts. 

vim ../../openfold/np/relax/amber_minimize.py # 修改 use_gpu 为 False