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

组件版本

hdk: 25.0.rc1.1
cann：8.3.RC1
python：3.10
torch：2.6.0
torch-npu：2.6.0

拉取模型代码

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

环境准备

  • 创建虚拟环境
conda create -n esm2 python=3.10 -y
conda activate esm2
  • 安装基础依赖
pip install torch==2.6.0 torch-npu==2.6.0 pyyaml numpy decorator attrs matplotlib
  • 验证 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 版本是否匹配

  • 安装fair-esm

pip3 install fair-esm
pip install "fair-esm[esmfold]"

推理

# source cann
source /usr/local/Ascend/ascend-toolkit/set_env.sh

# 指定可见 NPU 卡（0,1,2,3)
export ASCEND_RT_VISIBLE_DEVICES=0

# 运行
python pretrained.py
  • 查看生成的蛋白质序列接触图
ls -lh plot_*.png