​作者​：昇腾实战派
​知识地图​：https://blog.csdn.net/Lumos_Lovegood/article/details/161601003

本文记录蛋白质语言模型 ESM-2 迁移适配到华为昇腾（Ascend）NPU 平台的完整过程，涵盖模型原理、整体架构、环境准备、核心代码解析、推理运行与结果展示，力求给出一份可复现的实操指南。

一、ESM-2 介绍

ESM-2（Evolutionary Scale Modeling 2）是 Meta AI（FAIR）推出的蛋白质语言模型，它把自然语言处理里的掩码语言建模（Masked Language Modeling, MLM）思想搬到了蛋白质序列上：把一条氨基酸序列看作一句"蛋白质语言"，通过在海量蛋白序列（UniRef50）上做自监督预训练，让模型学会氨基酸之间的进化与结构约束。

和依赖多序列比对（MSA）的方法不同，ESM-2 只需要单条序列作为输入，就能输出蕴含丰富结构与功能信息的残基级表征（representation）。这些表征可以直接用于下游任务：

  • 接触图预测（residue-residue contact prediction）：无监督地从注意力图中读出残基之间的空间接触关系；
  • 结构预测：作为 ESMFold 的主干，端到端从序列直接预测三维结构；
  • 变异效应预测、二级结构、远程同源检测等蛋白质理解任务的特征提取器。

ESM-2 提供从 8M 到 15B 的一系列参数规模，方便在精度与算力之间权衡：

模型名称 层数 参数量 Embedding 维度 注意力头数
esm2_t6_8M_UR50D 6 8M 320 20
esm2_t12_35M_UR50D 12 35M 480 20
esm2_t30_150M_UR50D 30 150M 640 20
esm2_t33_650M_UR50D 33 650M 1280 20
esm2_t36_3B_UR50D 36 3B 2560 40
esm2_t48_15B_UR50D 48 15B 5120 40

注：模型命名规则为 esm2_t{层数}_{参数量}_UR50D，其中 UR50D 表示训练数据来自 UniRef50（Dense）。本文以最常用的 esm2_t33_650M_UR50D（33 层 / 650M 参数）为例。

二、整体架构

ESM-2 本质上是一个只含编码器（Encoder-only）的 Transformer，采用 RoBERTa 式的架构，主要由三部分组成：输入嵌入层、堆叠的 Transformer 编码层、以及面向不同任务的输出头（语言模型头 / 接触预测头）。整体数据流如下：

氨基酸序列
    │  Tokenize（Alphabet 词表映射）
    ▼
Token Embedding（+ Token Dropout 缩放）
    │
    ▼
┌─────────────────────────────┐
│  N × Transformer Layer      │
│   ├─ Multi-Head Attention   │  ← 旋转位置编码 RoPE
│   │    (Pre-LayerNorm)      │
│   └─ Feed-Forward Network   │  ← GELU 激活
│        (Pre-LayerNorm)      │
└─────────────────────────────┘
    │
    ├──► LM Head ────► logits（词表分布，用于 MLM）
    └──► Contact Head ► 接触图（残基接触概率矩阵）

对应源码见 esm/model/esm2.py，下面逐层拆解。

2.1 输入嵌入层

词表（Alphabet） 定义在 esm/constants.py 与 esm/data.py。标准氨基酸 token 共 25 个，加上特殊 token 组成完整词表：

# esm/constants.py
proteinseq_toks = {
    'toks': ['L', 'A', 'G', 'V', 'S', 'E', 'R', 'T', 'I', 'D', 'P', 'K', 'Q',
             'N', 'F', 'Y', 'M', 'H', 'W', 'C', 'X', 'B', 'U', 'Z', 'O', '.', '-']
}

ESM-2 使用 ESM-1b 架构的词表规则：在序列前添加 <cls>（起始，即 BOS）、序列后添加 <eos>（结束），并用 <pad> 补齐、<mask> 表示被遮盖的残基。词表构建时还会补足到 8 的整数倍以对齐硬件。

Token 经过 nn.Embedding 映射为 d d d 维向量。ESM-2 的一个关键细节是 Token Dropout：预训练时约 15% 的 token 会被替换成 <mask>（其中 80% 真正置为 mask），推理时为了让嵌入的数值分布与训练一致，需要对嵌入做缩放校正：

x = x ⋅ 1 − p train 1 − p observed x = x \cdot \frac{1 - p_{\text{train}}}{1 - p_{\text{observed}}} x=x⋅1−pobserved​1−ptrain​​

其中 p train = 0.15 × 0.8 = 0.12 p_{\text{train}} = 0.15 \times 0.8 = 0.12 ptrain​=0.15×0.8=0.12 是训练时的期望掩码比例， p observed p_{\text{observed}} pobserved​ 是当前输入里实际被 mask 的比例。对应实现：

# esm/model/esm2.py  forward()
x = self.embed_scale * self.embed_tokens(tokens)
if self.token_dropout:
    x.masked_fill_((tokens == self.mask_idx).unsqueeze(-1), 0.0)
    mask_ratio_train = 0.15 * 0.8
    src_lengths = (~padding_mask).sum(-1)
    mask_ratio_observed = (tokens == self.mask_idx).sum(-1).to(x.dtype) / src_lengths
    x = x * (1 - mask_ratio_train) / (1 - mask_ratio_observed)[:, None, None]

注：与 ESM-1b 不同，ESM-2 不使用可学习/正弦的绝对位置嵌入，位置信息完全交给下面的旋转位置编码（RoPE）在注意力内部注入。

2.2 Transformer 编码层

模型堆叠 num_layers 个 TransformerLayer（见 esm/modules.py），每层是标准的 Pre-LayerNorm 结构，包含多头自注意力和前馈网络两个子层，各自带残差连接：

# esm/modules.py  TransformerLayer.forward()
residual = x
x = self.self_attn_layer_norm(x)          # Pre-LN
x, attn = self.self_attn(query=x, key=x, value=x, ...)
x = residual + x                          # 残差

residual = x
x = self.final_layer_norm(x)              # Pre-LN
x = gelu(self.fc1(x))                     # 前馈 + GELU
x = self.fc2(x)
x = residual + x                          # 残差

其中：

  • 自注意力（Self-Attention）：查询 Q Q Q、键 K K K、值 V V V 均由同一输入线性投影得到，注意力权重按缩放点积计算：

Attention ( Q , K , V ) = softmax ( Q K ⊤ d k ) V \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V Attention(Q,K,V)=softmax(dk​ ​QK⊤​)V

其中 d k d_k dk​ 是每个注意力头的维度（head_dim = embed_dim / num_heads）， d k \sqrt{d_k} dk​ ​ 是缩放因子，防止点积过大导致 softmax 梯度消失。

  • 前馈网络（FFN）：两层线性变换，中间维度是 embedding 维度的 4 倍（4 * embed_dim），激活函数为 GELU：

FFN ( x ) = W 2 ⋅ GELU ( W 1 x + b 1 ) + b 2 \text{FFN}(x) = W_2 \cdot \text{GELU}(W_1 x + b_1) + b_2 FFN(x)=W2​⋅GELU(W1​x+b1​)+b2​

GELU（高斯误差线性单元）的定义为 GELU ( x ) = x ⋅ Φ ( x ) \text{GELU}(x) = x \cdot \Phi(x) GELU(x)=x⋅Φ(x)，其中 Φ ( x ) \Phi(x) Φ(x) 是标准正态分布的累积分布函数，源码用误差函数 erf 精确实现：

# esm/modules.py
def gelu(x):
    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))

2.3 旋转位置编码（RoPE）

ESM-2 用旋转位置编码（Rotary Position Embedding）替代绝对位置嵌入，实现见 esm/rotary_embedding.py。RoPE 的核心思想是：通过一个只依赖位置的旋转矩阵作用到 Q Q Q、 K K K 上，使得注意力点积天然携带相对位置信息。

对位置 m m m 处的向量，旋转所用的频率为：

θ i = 10000 − 2 i / d , i = 0 , 1 , … , d / 2 − 1 \theta_i = 10000^{-2i/d}, \quad i = 0, 1, \dots, d/2 - 1 θi​=10000−2i/d,i=0,1,…,d/2−1

具体实现是把向量后半段"旋转"后与前半段拼接，再与 cos ⁡ \cos cos、 sin ⁡ \sin sin 表逐元素相乘：

# esm/rotary_embedding.py
def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(x, cos, sin):
    cos = cos[:, : x.shape[-2], :]
    sin = sin[:, : x.shape[-2], :]
    return (x * cos) + (rotate_half(x) * sin)

RoPE 的优势在于：无需额外可学习参数、能自然外推到训练时未见过的序列长度、且相对位置关系在点积中显式保留，这对蛋白质这类长序列建模尤为重要。

2.4 输出头

经过所有编码层并做最终 LayerNorm 后，隐藏表征分别送入两个输出头：

（1）语言模型头（RobertaLMHead） —— 把每个位置的表征投影回词表维度，输出各氨基酸的概率分布，用于掩码预测：

# esm/modules.py  RobertaLMHead.forward()
x = self.dense(features)
x = gelu(x)
x = self.layer_norm(x)
x = F.linear(x, self.weight) + self.bias   # 投影回词表大小

（2）接触预测头（ContactPredictionHead） —— 这是 ESM-2 最具特色的能力。它把所有层、所有注意力头的注意力权重（共 num_layers × attention_heads 个通道）堆叠起来，依次做对称化（symmetrize）、平均乘积校正（APC, Average Product Correction），最后用一个逻辑回归 + Sigmoid 输出每对残基接触的概率：

APC ( x i j ) = x i j − x i ⋅   x ⋅ j x ⋅ ⋅ \text{APC}(x_{ij}) = x_{ij} - \frac{x_{i \cdot}\, x_{\cdot j}}{x_{\cdot \cdot}} APC(xij​)=xij​−x⋅⋅​xi⋅​x⋅j​​

其中 x i ⋅ x_{i\cdot} xi⋅​、 x ⋅ j x_{\cdot j} x⋅j​、 x ⋅ ⋅ x_{\cdot\cdot} x⋅⋅​ 分别是第 i i i 行之和、第 j j j 列之和、全矩阵之和。APC 用于扣除背景噪声，是共进化分析里的经典技巧。

# esm/modules.py  ContactPredictionHead
def symmetrize(x):
    return x + x.transpose(-1, -2)          # 对称化

def apc(x):
    a1 = x.sum(-1, keepdims=True)
    a2 = x.sum(-2, keepdims=True)
    a12 = x.sum((-1, -2), keepdims=True)
    avg = a1 * a2
    avg.div_(a12)
    return x - avg                          # 平均乘积校正

注：接触预测依赖单独的回归权重（-contact-regression.pt）。加载模型时若缺少这份权重，接触图结果将不准确，程序会给出告警。

三、实验环境

3.1 硬件与组件版本

本文在昇腾 NPU 环境下完成推理，各组件版本如下：

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

注：torch-npu 的版本需与 torch 主版本严格对应（此处均为 2.6.0），同时要与已安装的 CANN 版本匹配，否则会在导入或算子下发时报错。

3.2 拉取模型代码

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

3.3 创建虚拟环境

conda create -n esm2 python=3.10 -y
conda activate esm2

3.4 安装基础依赖

pip install torch==2.6.0 torch-npu==2.6.0 pyyaml numpy decorator attrs matplotlib
  • torch / torch-npu：深度学习框架及昇腾适配层；
  • matplotlib：绘制接触图；
  • decorator、attrs、pyyaml、numpy：torch_npu 的运行时依赖。

3.5 验证 PyTorch 与 torch_npu 安装

执行以下命令，检查昇腾环境是否可用：

python3 -c "import torch;import torch_npu; a = torch.randn(3, 4).npu(); print(a + a);"

若输出类似下面的张量（device='npu:0'），说明安装成功：

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（CANN 环境变量未生效）；
  • decorator 等运行时依赖是否已安装；
  • CANN 与 torch_npu 版本是否匹配。

3.6 安装 fair-esm

ESM 官方以 fair-esm 包形式发布，安装它即可通过 import esm 使用：

pip3 install fair-esm
pip install "fair-esm[esmfold]"     # 如需 ESMFold 结构预测能力

四、推理代码解析（pretrained.py）

推理入口脚本为项目根目录下的 pretrained.py，它完整演示了"加载模型 → 准备数据 → 前向推理 → 提取表征 → 绘制接触图"的流程。下面分段解析。

4.1 设备选择与迁移

import torch
import torch_npu
import esm

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

关键点在于 import torch_npu：它把 NPU 注册为 PyTorch 的后端设备，之后就能像用 cuda 一样用 npu。后续通过 .to(device) 把模型和数据搬到 NPU。

4.2 加载预训练模型

model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
model = model.to(device)

batch_converter = alphabet.get_batch_converter()
model.eval()   # 关闭 dropout，保证结果确定
  • esm.pretrained.esm2_t33_650M_UR50D() 返回 (模型, 词表) 二元组。首次调用会自动从官方地址下载权重并缓存到 ~/.cache/torch/hub/checkpoints/；
  • alphabet.get_batch_converter() 返回一个批处理转换器，负责把"序列字符串"变成"模型可读的 token 张量"；
  • model.eval() 切换到评估模式。

注：esm2_t33_650M_UR50D 权重约 2.5GB，加上接触回归权重需联网下载。离线环境可提前下载 .pt 文件，用 esm.pretrained.load_model_and_alphabet_local("路径/xxx.pt") 加载，并确保回归权重同目录。

4.3 准备输入数据

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)

输入是 (名称, 序列) 的列表。注意第三、四条序列里的 <mask>：这演示了 MLM 能力——模型会去预测被遮盖位置最可能的氨基酸。batch_converter 会自动补上 <cls>、<eos> 并按最长序列做 padding。batch_lens 记录每条序列的真实长度（排除 padding），用于后续裁剪。

4.4 前向推理

with torch.no_grad():
    results = model(batch_tokens, repr_layers=[33], return_contacts=True)
token_representations = results["representations"][33]
  • torch.no_grad()：推理不需要梯度，节省显存并加速；
  • repr_layers=[33]：请求返回第 33 层（最后一层）的隐藏表征；
  • return_contacts=True：同时计算接触图。

返回的 results 是一个字典，包含 logits（词表分布）、representations（各层表征）、attentions（注意力权重）、contacts（接触图）。

4.5 提取序列级表征

sequence_representations = []
for i, tokens_len in enumerate(batch_lens):
    sequence_representations.append(token_representations[i, 1 : tokens_len - 1].mean(0))

模型输出的是残基级（per-token）表征。要得到一条序列的整体表征，常用做法是对所有残基位置取平均。这里 1 : tokens_len - 1 的切片刻意跳过了首位的 <cls> 和末位的 <eos>，只对真实氨基酸求均值。

4.6 绘制接触图

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")

results["contacts"] 是每条序列的残基接触概率矩阵（ L × L L \times L L×L， L L L 为序列长度）。矩阵中 ( i , j ) (i, j) (i,j) 位置的值越大，代表残基 i i i 和 j j j 在三维空间中接触的概率越高。用 matshow 把矩阵可视化成热力图并保存为 PNG。注意绘图前要 .cpu() 把张量从 NPU 搬回 CPU。

五、运行推理

5.1 执行脚本

# 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

5.2 预期日志

脚本在各阶段都打印了进度，正常运行时终端输出类似：

using device:npu
Loading model...
Model loaded successfully.
Moving model to device...
Model moved to device.
Preparing data...
Data prepared.
Tokens moved to device.
Starting forward pass...
Forward pass completed.
Token representations extracted.
Generating sequence representations...
Sequence representations generated.
Plots saved successfully.

六、结果展示

6.1 查看生成的接触图

推理结束后，当前目录会生成每条输入序列对应的接触图：

ls -lh plot_*.png

将得到 plot_0.png、plot_1.png、plot_2.png、plot_3.png 四张图，分别对应输入的四条序列。

6.2 如何解读接触图

  • 图像是一个对称的方形热力图，横纵坐标都是残基序号；
  • 颜色越亮（值越接近 1）表示对应两个残基空间接触概率越高；
  • 对角线附近的高值是相邻残基的局部接触，属正常现象；远离对角线的亮点则往往对应二级结构（如 β 折叠）或长程折叠接触，是判断蛋白质结构的重要线索。

参考资料

# 语言模型 # 人工智能 # 自然语言处理 # python # pytorch
Logo
鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。

更多推荐

华为昇腾910B拉起MinerU2.5-Pro-2605-1.2B示例

昇腾910B拉起模型MinerU2.5-Pro-2605-1.2B

avatar 鲲鹏昇腾开发者社区

像管理流水线一样管理研发：Gitee 用“七大车间“重构软件生产模式

类别 适配清单 | 国产芯片 | 鲲鹏、飞腾、海光、龙芯 | | 国产操作系统 | 麒麟、OpenEuler、统信 UOS | | 国产数据库 | OceanBase、达梦、TDSQL、TiDB、GaussDB | | 国产中间件 | TongWeb、TAS、RocketMQ、东方通中间件 | | 部署模式 | 单机/集群/K8s 容器化/跨地域多机房/私有化 |Gitee 的解法：软件工厂将研发

avatar 鲲鹏昇腾开发者社区

图像算法模型NPU适配与算法服务实战指南

本项目提供了一个开箱即用的图像推理服务框架，支持主流目标检测模型（如YOLO系列、MMDetection等），可适配NPU/CUDA/CPU多种硬件环境。核心功能包括：多模型多后端支持、昇腾NPU深度适配（自动选卡、兼容性补丁、性能优化）、高并发批处理设计、生产级健壮性保障。通过统一的REST API提供推理能力，用户只需配置模型文件即可快速部署服务。技术架构基于FastAPI实现，包含模型管理、

avatar 鲲鹏昇腾开发者社区