lifescience on ascend

ascend life science 相关仓库¶

Warning

以下所有项目，如无特殊说明，均在 AI-01 节点运行，使用前向管理员申请 AI 队列使用权限，并在群共享表格超算GPU-NPU使用登记中登记使用情况。

xfold¶

evo2¶

项目核心定位¶

Evo2 是一款先进的基因组建模与设计工具，基于前沿深度学习架构（StripedHyena 2）构建，专注于 DNA、RNA 及蛋白质的序列预测与设计任务，支持单核苷酸分辨率建模，最大上下文窗口可达 100 万碱基对，是生物信息学研究与实践的高效辅助工具。

核心功能¶

序列生成：通过 DNA 序列或物种提示，生成目标序列并标注编码区；
序列评分：对输入 DNA 序列进行困惑度与核苷酸熵值评估；
辅助可视化：支持原核生物序列的 3D 蛋白质可视化（基于 ESMFold）。

链接¶

线上试用地址：Evo Designer（无需部署，直接体验序列生成与评分）
项目开源地址：GitHub - ArcInstitute/evo2（含完整代码、示例笔记本、本地部署教程）
昇腾版项目地址: 待补充

使用¶

# 登录节点下载适配好的的evo2镜像 (已下载则忽略此步骤)
$ wget https://obs-whaicc-fae-public.obs.cn-central-221.ovaijisuan.com:443/singularity-img/ai4s/singularity-evo2.03.tar
# 解压
$ tar -xvf singularity-evo2.03.tar && cd singularity-evo2.03

# 交互模式进入ai节点
$ dsub -R 'npu=1,cpu=4'  -aa -q AI   -I bash
# 载入 singularity
$ module load arm/singularity/4.1.5 
# 测试运行
$ bash run.sh

默认执行脚本内容为使用 evo2_7b 模型，根据提示词，生成基因序列。

esm2¶

原版项目地址 https://github.com/facebookresearch/esm

昇腾版项目地址 https://modelers.cn/models/Ascend-AI4S/ESM2

自行安装环境¶

# 交互模式进入ai节点
$ dsub -R 'npu=1,cpu=4'  -aa -q AI   -I bash

# 载入 npu 版 pytorch
$ module load arm/python/3.11.11 arm/torch_npu/2.3.1.post4-python311

# 安装基础依赖，默认安装路径为 ~/.local/lib/python3.11/site-packages/
$ pip3.11 install pyyaml decorator attrs matplotlib

# 安装 ESM 推理模块
$ pip3.11 install fair-esm
$ pip3.11 install "fair-esm[esmfold]"

$ pip3.11 install torch_npu/2.3.1.post4-python311
$ pip3.11 install "numpy<2.0"

使用集群环境¶

# 交互模式进入ai节点
$ dsub -R 'npu=1,cpu=4'  -aa -q AI   -I bash
$ module load arm/esm2/1.0-py3.11

测试¶

示例1¶

提取蛋白质序列的残基级与整链级特征，并画出自注意力预测的残基接触图。

pretrain.py

import torch
import torch_npu
import esm

# use ascend npu
device = torch.device('npu' if torch_npu.npu.is_available() else 'cpu')
print(f"using device:{device}")

# Load ESM-2 model
model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()

model = model.to(device)

batch_converter = alphabet.get_batch_converter()
model.eval()  # disables dropout for deterministic results

# Prepare data (first 2 sequences from ESMStructuralSplitDataset superfamily / 4)
data = [
    ("protein1", "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"),
    ("protein2", "KALTARQQEVFDLIRDHISQTGMPPTRAEIAQRLGFRSPNAAEEHLKALARKGVIEIVSGASRGIRLLQEE"),
    ("protein2 with mask","KALTARQQEVFDLIRD<mask>ISQTGMPPTRAEIAQRLGFRSPNAAEEHLKALARKGVIEIVSGASRGIRLLQEE"),
    ("protein3",  "K A <mask> I S Q"),
]
batch_labels, batch_strs, batch_tokens = batch_converter(data)

batch_tokens = batch_tokens.to(device)

batch_lens = (batch_tokens != alphabet.padding_idx).sum(1)

# Extract per-residue representations (on CPU)
with torch.no_grad():
    results = model(batch_tokens, repr_layers=[33], return_contacts=True)
token_representations = results["representations"][33]

# Generate per-sequence representations via averaging
# NOTE: token 0 is always a beginning-of-sequence token, so the first residue is token 1.
sequence_representations = []
for i, tokens_len in enumerate(batch_lens):
    sequence_representations.append(token_representations[i, 1 : tokens_len - 1].mean(0))

# Look at the unsupervised self-attention map contact predictions
import matplotlib.pyplot as plt
for idx,((_, seq), tokens_len, attention_contacts) in enumerate(zip(data, batch_lens, results["contacts"])):
    plt.figure()
    plt.matshow(attention_contacts[: tokens_len, : tokens_len].cpu().numpy())
    plt.title(seq)
    plt.savefig(f"plot_{idx}.png")

运行 python pretrain.py，当前目录下会生成几个 png 文件。

示例2¶

提取蛋白质序列特征

feature.py

import torch
import torch_npu

import esm

device = torch.device('npu' if torch_npu.npu.is_available() else 'cpu')
#device = torch.device('npu')
print(f"using device:{device}")

# 加载预训练模型
model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
model = model.to(device)
model.eval()                        # 推理模式

batch_converter = alphabet.get_batch_converter()

# 准备蛋白质序列数据
data = [("protein1", "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG")]
batch_labels, batch_strs, batch_tokens = batch_converter(data)
batch_tokens = batch_tokens.to(device)      # <= 输入也要 to(device)

# 提取序列特征
with torch.no_grad():
    results = model(batch_tokens, repr_layers=[33])
token_representations = results["representations"][33]

print(f"提取的特征形状: {token_representations.shape}")

运行 python feature.py。

模型说明¶

ESM-2 一共有多个版本，主要区别在于：层数（depth）、参数量（size）、推理速度和精度权衡。这些版本都遵循相同的 Transformer 编码器架构，只是在大小和计算能力上有差异。

模型命名：esm2_t<层数>_<参数量>_UR50D

t<层数>：比如 t12 表示有 12 层 Transformer 编码器。
<参数量>：大概的参数数量，比如 35M, 3B 等。
UR50D：代表训练用的数据集（Uniref50），是去冗余后的蛋白质数据库。

模型名称（Hugging Face 名）	层数	参数量	说明
esm2_t6_8M_UR50D	6	8M	极小模型，适合快速原型，快速测试 / 教学 / CPU 调用
esm2_t12_35M_UR50D	12	35M	中等小型，推荐用于入门任务，快速测试 / 教学 / CPU 调用
esm2_t30_150M_UR50D	30	150M	中等模型，效果与效率平衡，标准下游任务建模
esm2_t33_650M_UR50D	33	650M	较大模型，适合更复杂任务，结构预测 / 蛋白功能预测
esm2_t36_3B_UR50D	36	3B	超大模型，强大建模能力，结构预测 / 蛋白功能预测
esm2_t48_15B_UR50D	48	15B	最大模型，性能最强但最重，结构预测 / 蛋白功能预测

参考资料¶

AI模型学习-ESM2

官方 examples

基于蛋白质语言模型的蛋白质相互作用(PPI)的预测

蛋白质语言模型微调cookbook

使用 ESM 计算蛋白序列的突变影响分数

esmfold¶

项目地址 https://ai.gitcode.com/zzhihaoo/ESMfold

使用¶

# 交互模式进入ai节点
$ dsub -R 'npu=1,cpu=4'  -aa -q AI   -I bash
$ module load arm/singularity/4.1.5
$ singularity exec -B /usr/local/Ascend/driver:/usr/local/Ascend/driver:ro   -B /usr/local/Ascend/ascend-toolkit/8.2.RC1/opp:/usr/local/Ascend/opp:ro /share/software/image/arm/esmfold.2.0.1.sif esm-fold -i test.fasta -o .

test.fasta

>test
MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG

esm3¶

LigandMPNN¶

项目地址 https://gitcode.com/AI4Science/LifeScience/blob/main/PyTorch/LigandMPNN

自行安装环境¶

# 交互模式进入ai节点
$ dsub -R 'npu=1,cpu=4'  -aa -q AI   -I bash

# 默认安装路径 ~/.local/lib/python3.11/site-packages/
$ module load arm/python/3.11.11 
# 安装torch-npu
$ pip3.11 install  torch-npu==2.3.1.post4
# 拉取模型代码
git clone https://gitcode.com/AI4Science/LifeScience.git
cd LifeScience/PyTorch/LigandMPNN
git clone https://github.com/dauparas/LigandMPNN.git && cd LigandMPNN
git checkout 26ec57ac976ade5379920dbd43c7f97a91cf82de
git apply ../patch/LigandMPNN.patch
# 安装环境依赖
$ pip3.11 install -r requirements.txt

获取模型权重 bash get_model_params.sh "./model_params"。

测试数据目录 inputs。

推理 bash run_examples.sh。

使用集群环境¶

使用集群预装好的环境

# 交互模式进入ai节点
$ dsub -R 'npu=1,cpu=4'  -aa -q AI   -I bash

# 下载推理代码
git clone https://gitcode.com/AI4Science/LifeScience.git
cd LifeScience/PyTorch/LigandMPNN

# 加载环境
$ module load arm/ligandmpnn/1.0-py3.11

# 下载模型权重
$ bash get_model_params.sh "./model_params"

# 推理，数据在 inputs 下
$ bash run_examples.sh

BioGPT¶

项目地址 https://modelers.cn/models/Ascend-AI4S/BioGPT

使用集群环境¶

# 交互模式进入ai节点
$ dsub -R 'npu=1,cpu=4'  -aa -q AI   -I bash
# 将运行脚本拷贝至自己的目录下
$ cp /share/software/image/arm/singularity-biogpt-02/run.sh .
# 运行
$ sh  run.sh  "Immunotherapy for breast cancer"
...
输入文本：Immunotherapy for breast cancer
开始生成（最大长度：1024，束宽：5）

==================================================
生成结果：
==================================================
Immunotherapy for breast cancer: current status and future prospects for the treatment of metastatic disease. (1) Breast Cancer Immunotherapy: Current Status and Future Prospects for The Treatment of Metastatic Disease: (2) The Role of Immunotherapy in the Management of Early-Stage and Locally Advanced Breast Cancers: A Systematic Review and Meta-Analysis of the Literature (3) Immunotherapy as an Adjuvant Therapy for Patients with High-Risk, Hormone Receptor-Positive, Human Epidermal Growth Factor Receptor 2 (HER2) -Negative, Estrogen Receptor (ER) -negative, and Her2 / neu-Overexpressing (HR + / HER2-) Invasive Breast Carcinomas.
==================================================

RFdiffusion¶

原项目地址 https://github.com/RosettaCommons/RFdiffusion

昇腾项目地址 https://modelers.cn/models/Ascend-AI4S/RFdiffusion

使用集群环境¶

# 交互模式进入ai节点
$ dsub -R 'npu=1,cpu=4'  -aa -q AI   -I bash
$ module load arm/rfdiffsion/ascend
$ ln -s    /share/software/app/arm/rfdiffsion_pyg_ascend/{examples,ops,rfdiffusion,scripts,config,models}  .
# 测试运行
$ sh examples/design_unconditional.sh

ProteinMPNN¶

项目地址 https://modelers.cn/models/Ascend-AI4S/ProteinMPNN

自行安装环境¶

# 交互模式进入ai节点
$ dsub -R 'npu=1,cpu=4'  -aa -q AI   -I bash
# 安装 torch-npu 及相关依赖
$ module load arm/python/3.11.11
$ pip3.11 install torch-npu==2.6.0
$ pip3.11 install "numpy<2.0" pyyaml 

# 安装mx_driving，用于训练，推理可以不要
$ pip3.11 install wheel setuptools decorator scipy attrs psutil cloudpickle synr==0.5.0 tornado protobuf==4.25.8 
$ git clone https://gitee.com/ascend/DrivingSDK.git
$ cd DrivingSDK
$ source /usr/local/Ascend/ascend-toolkit/set_env.sh
$ export CPATH="$HOME/.local/lib/python3.11/site-packages/torch_npu/include/"
$ export LIBRARY_PATH="$HOME/.local/lib/python3.11/site-packages/torch_npu/lib:$LIBRARY_PATH"
$ bash ci/build.sh --python=3.11
$ pip3.11 install --force-reinstall dist/mx_driving-1.0.0+git3e3c655-cp311-cp311-linux_aarch64.whl 

# 下载 ProteinMPNN
$ git clone https://modelers.cn/Ascend-AI4S/ProteinMPNN.git

使用集群环境¶

# 交互模式进入ai节点
$ dsub -R 'npu=1,cpu=4'  -aa -q AI   -I bash
$ module load arm/proteinmpnn/1.0-py3.11

推理/训练¶

下载 ProteinMPNN 相关代码和权重

$ git clone https://modelers.cn/Ascend-AI4S/ProteinMPNN.git

推理

$ cd ProteinMPNN/training/
$ bash test_inference.sh

训练

# 下载数据集
$ wget https://files.ipd.uw.edu/pub/training_sets/pdb_2021aug02.tar.gz
# 解压数据集
$ tar -xzf pdb_2021aug02.tar.gz
$ cd ProteinMPNN/training/
$ bash submit_exp_020.sh

报错¶

重装 torch_npu¶

报错

$ python test.py
Traceback (most recent call last):
  File "/share/home/software/package/esm2/pretrained.py", line 2, in <module>
    import torch_npu
  File "/share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch_npu/__init__.py", line 12, in <module>
    from torch.distributed.fsdp import sharded_grad_scaler
  File "/share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch/distributed/fsdp/__init__.py", line 1, in <module>
    from ._flat_param import FlatParameter as FlatParameter
  File "/share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch/distributed/fsdp/_flat_param.py", line 30, in <module>
    from torch.distributed.fsdp._common_utils import (
  File "/share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch/distributed/fsdp/_common_utils.py", line 35, in <module>
    from torch.distributed.fsdp._fsdp_extensions import FSDPExtensions
  File "/share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch/distributed/fsdp/_fsdp_extensions.py", line 8, in <module>
    from torch.distributed._tensor import DeviceMesh, DTensor
  File "/share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch/distributed/_tensor/__init__.py", line 6, in <module>
    import torch.distributed._tensor.ops
  File "/share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch/distributed/_tensor/ops/__init__.py", line 2, in <module>
    from .embedding_ops import *  # noqa: F403
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch/distributed/_tensor/ops/embedding_ops.py", line 8, in <module>
    import torch.distributed._functional_collectives as funcol
  File "/share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch/distributed/_functional_collectives.py", line 12, in <module>
    from . import _functional_collectives_impl as fun_col_impl
  File "/share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch/distributed/_functional_collectives_impl.py", line 36, in <module>
    from torch._dynamo import assume_constant_result
  File "/share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch/_dynamo/__init__.py", line 2, in <module>
    from . import convert_frame, eval_frame, resume_execution
  File "/share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 40, in <module>
    from . import config, exc, trace_rules
  File "/share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch/_dynamo/trace_rules.py", line 50, in <module>
    from .variables import (
  File "/share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch/_dynamo/variables/__init__.py", line 34, in <module>
    from .higher_order_ops import (
  File "/share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch/_dynamo/variables/higher_order_ops.py", line 13, in <module>
    import torch.onnx.operators
  File "/share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch/onnx/__init__.py", line 46, in <module>
    from ._internal.exporter import (  # usort:skip. needs to be last to avoid circular import
ImportError: cannot import name 'DiagnosticOptions' from 'torch.onnx._internal.exporter' (/share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch/onnx/_internal/exporter/__init__.py)

解决

$ pip3.11 uninstall -y torch torch_npu
$ rm -rf /share/software/app/arm/esm2/1.0/lib/python3.11/site-packages/torch*
$ pip3.11 install torch-npu==2.3.1.post4

程序被终止¶

交互模式下(dsub -R 'npu=1,cpu=4' -aa -q AI -I bash) 运行 AI 相关程序，如果程序运行一段时间后突然被终止并被退出了 AI 节点，同时出现如下 killed 的提示，则可能是程序内存(非NPU显存)使用较大，超过了作业能使用的内存限制，建议调高作业核数，如 dsub -R 'npu=1,cpu=10' -aa -q AI -I bash，多瑙作业内存限制见多瑙使用-内存限制。

Killed
[username@AI-01 esmfold]

本文阅读量次
本站总访问量次