AI4S/ESMfold
模型介绍文件和版本Pull Requests讨论分析

组件版本

hdk：25.5.0
cann：3.RC1
python：3.10
torch：2.6.0
torch_npu：2.6.0

下载源码

git clone https://atomgit.com/AI4Science/ESMfold.git
cd ESMfold

环境准备

1.创建虚拟环境

conda create -n esmfold python=3.10 -y
conda activate esmfold
  1. 安装依赖
# 下载并安装PyTorch框架
wget https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp310-cp310-manylinux_2_28_aarch64.whl
pip3 install torch-2.6.0+cpu-cp310-cp310-manylinux_2_28_aarch64.whl

# 下载并安装torch_npu插件
wget https://gitcode.com/Ascend/pytorch/releases/download/v7.3.0-pytorch2.6.0/torch_npu-2.6.0.post5-cp310-cp310-manylinux_2_28_aarch64.whl
pip3 install torch_npu-2.6.0.post5-cp310-cp310-manylinux_2_28_aarch64.whl

# 安装项目基础依赖包 
pip install pyyaml decorator attrs psutil scipy setuptools==65.5.1 numpy==1.23.5
pip install --upgrade deepspeed
  1. 验证 PyTorch 与 torch_npu 安装 执行以下命令检查安装是否成功：
python3 -c "import torch;import torch_npu; a = torch.randn(3, 4).npu(); print(a + a);"

若输出类似以下信息，说明安装成功：

tensor([[-0.6066,  6.3385,  0.0379,  3.3356],
        [ 2.9243,  3.3134, -1.5465,  0.1916],
        [-2.1807,  0.2008, -1.1431,  2.1523]], device='npu:0')

若报错，排查顺序：

  • set_env.sh 是否已 source
  • decorator 等运行时依赖是否安装
  • CANN 与 torch_npu 版本是否匹配

安装esmfold

pip install fair-esm
pip install "fair-esm[esmfold]"
pip install -e .

安装openfold

git clone https://atomgit.com/AI4Science/openfold-1.0.0.git
cd openfold-1.0.0

export ASCEND_HOME_PATH=/usr/local/Ascend/ascend-toolkit/
export TORCH_DEVICE_BACKEND_AUTOLOAD=0

pip install -e . --no-build-isolation
pip install pytorch-lightning==1.9.5
pip install git+https://github.com/NVIDIA/dllogger.git
cd ..

下载权重

wget https://dl.fbaipublicfiles.com/fair-esm/models/esmfold_3B_v1.pt
wget https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t36_3B_UR50D.pt
wget https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t36_3B_UR50D-contact-regression.pt
将权重移动至/root/.cache/torch/hub/checkpoints/
mkdir -p /root/.cache/torch/hub/checkpoints/
mv *.pt /root/.cache/torch/hub/checkpoints/

推理

source /usr/local/Ascend/ascend-toolkit/set_env.sh
# 指定可见 NPU 卡（0,1,2,3)
export ASCEND_RT_VISIBLE_DEVICES=0

esm-fold -i test.fasta -o ./
  • 如遇报错ModuleNotFoundError: No module named 'torch._six'
vi /root/miniconda3/envs/esmfold/lib/python3.10/site-packages/deepspeed/runtime/utils.py
vi /root/miniconda3/envs/esmfold/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py

将 from torch._six import inf 替换为 from math import inf

Docker镜像

已预置 Conda 运行环境、模型权重、ESMfold源码，开箱即用，可快速部署运行。

# 拉取镜像
sudo docker pull swr.cn-north-4.myhuaweicloud.com/ascend_ai4s/esmfold:v1
  • 项目源码路径：/root/ESMfold/
  • 已配置好 Conda 虚拟环境
  • 已内置模型权重文件，无需额外下载