作者：昇腾实战派

知识地图：https://blog.csdn.net/Lumos_Lovegood/article/details/161601003

背景概述

自监督深度语言建模在自然语言处理领域取得了前所未有的成功，近年来也被广泛应用于生物序列分析。然而，现有的蛋白质语言模型（如 ESM、TAPE-Transformer、ProtTrans 等）大多从自然语言领域直接引入现有架构和预训练任务，未能充分利用蛋白质的独特特性——蛋白质没有明确的多字母构建块（如单词和句子），长度变化远大于句子，且三维结构导致远距离位置之间存在大量相互作用。

ProteinBERT 由耶路撒冷希伯来大学 Nadav Brandes 团队于 2022 年发表于 Bioinformatics，是一种专门为蛋白质设计的深度语言模型。它的核心创新在于：引入了局部（序列级）与全局（蛋白质级）双重表示机制，预训练方案结合了掩码语言建模与基因本体论（GO）注释预测双任务，且采用线性复杂度的全局注意力层替代传统自注意力，使模型能高效处理任意长度的蛋白质序列。整个模型仅约 1600 万参数，比其他蛋白质语言模型（通常 3800 万以上）小得多、快得多，却在多个基准上达到接近甚至超越 SOTA 的性能。

本文介绍 ProteinBERT 模型的 PyTorch + 昇腾 Ascend NPU 适配版本——从原始 TensorFlow/Keras 项目中转换预训练权重，在昇腾 AI 平台上完成推理与下游 benchmark 微调，为蛋白质功能预测和序列分析提供高效、可复现的技术方案。

模型介绍

ProteinBERT 概述

ProteinBERT 在约 1.06 亿条 UniRef90 蛋白质序列上进行自监督预训练，预训练任务包括：

• 掩码语言建模（MLM）：以 5% 概率随机替换氨基酸 token，模型需恢复原始序列
• GO 注释预测：随机删除/添加 GO 注释，模型需恢复已知注释

模型支持端到端处理局部输入输出（序列级别，如二级结构、信号肽位置）和全局输入输出（蛋白质级别，如荧光、稳定性、远程同源性分类）。

整体架构

ProteinBERT 的架构灵感来自 BERT，但有几项关键创新：

双路径表示

路径	张量形状	维度	作用
局部表示（Local）	B × L × d_local	d_local=128	捕获残基级别的序列特征
全局表示（Global）	B × d_global	d_global=512	捕获蛋白质整体功能特征

模型组成

1. 输入层

   • 序列输入：26 个 token（20 种标准氨基酸 + U + X + OTHER + START/END/PAD）通过 Embedding 层转为局部表示
   • GO 注释输入：8943 维 one-hot 向量通过全连接层转为全局表示

2. 核心主干：6 个 ProteinBERT Block

   每个 Block 包含：

   • 局部路径：窄卷积（kernel=9, dilation=1）+ 宽卷积（kernel=9, dilation=5）+ 位置全连接层 + LayerNorm
   • 全局路径：两层全连接 + LayerNorm
   • 局部→全局：全局注意力层（Global Attention），全局表示作为 Query，局部表示作为 Key/Value
   • 全局→局部：广播全连接层，将全局信息投影到每个序列位置

3. 输出层

   • 局部输出：序列每个位置的 token 概率（Softmax）
   • 全局输出：8943 个 GO 注释的预测概率（Sigmoid）

全局注意力层（Global Attention）

这是 ProteinBERT 的核心创新。传统自注意力复杂度为 O(L²)，而全局注意力的复杂度仅为 O(L)——全局表示（512维向量）作为 Query 与序列所有位置计算注意力，无需序列位置之间两两计算。这使得模型能高效处理极长的蛋白质序列（上万氨基酸），同时保持了全局信息与局部特征的双向流动。

具体实现：

• 4 个注意力头，Key 维度 64
• Query 来自全局表示，Key/Value 来自局部序列表示
• 使用 tanh 激活 Query/Key，GELU 激活 Value
• 注意力输出拼接后作为全局表示的残差更新

预训练与微调策略

预训练：

• 数据：UniRef90 约 1.06 亿蛋白质序列 + GO 注释
• 损失函数：序列分类交叉熵 + GO 注释二元交叉熵
• 训练约 28 天，~6.4 epochs，定期切换序列编码长度（128/512/1024）

微调协议：

1. 冻结所有预训练层，仅训练新增全连接层 40 epochs
2. 解冻所有层，继续训练 40 epochs
3. 在更大序列长度上做最终 epoch

应用场景

• 信号肽识别：预测蛋白质 N 端是否存在信号肽及切割位点
• 蛋白质稳定性预测：预测突变对蛋白质热稳定性的影响
• 荧光强度预测：预测 GFP 突变体的荧光活性
• 远程同源性检测：识别序列相似性极低但结构/功能相关的蛋白质
• 二级结构预测：预测每个残基的二级结构类型（螺旋/折叠/卷曲）
• 翻译后修饰预测：预测磷酸化等修饰位点

PyTorch + 昇腾 NPU 适配版本

迁移动机

原始 ProteinBERT 基于 Keras/TensorFlow 实现。为在昇腾 Ascend NPU 上高效运行推理与微调，本项目将核心模型代码转换为 PyTorch 实现，并通过 torch_npu 适配昇腾硬件加速。

仓库结构

proteinbert_pytorch/
├── README.md
├── setup.py
├── bin/
│   ├── convert_tf_to_pytorch       # 权重转换入口
│   ├── env_npu.sh                  # 昇腾环境配置
│   ├── inference_proteinbert_npu   # 推理入口
│   └── finetune_proteinbert_npu    # 微调入口
├── scripts/
│   ├── demo_scripts/               # 完整 benchmark 演示
│   │   ├── demo1_signalP_npu.py
│   │   ├── demo2_all_benchmarks_npu.py
│   │   └── demo3_attention_npu.py
│   ├── deploy_toolkit/             # 部署相关脚本
│   └── tools/                      # 调试与对比工具
├── protein_benchmarks/             # 下游 benchmark 数据集
└── proteinbert/
    ├── __init__.py
    ├── model.py                    # ProteinBERTModel PyTorch 实现
    ├── convert_weights.py          # TF pkl → PyTorch state_dict
    ├── inference.py                # 推理逻辑
    └── finetune.py                 # 微调逻辑

核心实现

ProteinBERTModel：包含序列 Embedding、全局输入 Dense、6 个 ProteinBERTBlock、序列输出 Dense 和注释输出 Dense。

ProteinBERTBlock：每个 Block 实现局部路径（窄卷积 + 宽卷积 + 全连接 + LayerNorm）和全局路径（全连接 + GlobalAttention + LayerNorm），以及两路径之间的信息交换。

GlobalAttention：多头全局注意力，使用 einsum 高效计算 Query（全局）与 Key/Value（序列）的注意力，复杂度为 O(L)。

当前覆盖能力

• 加载原始 TensorFlow pkl 预训练权重
• 转换为 PyTorch state_dict
• 提取 embedding（局部 + 全局）
• 在 CPU / Ascend NPU 上做推理
• 在 9 个下游 benchmark 上做微调与评估

当前版本不包含原始仓库完整的 UniRef 预训练流水线的 PyTorch 复刻。

版本信息

软件	版本
CANN	8.2+
Python	3.10
PyTorch	2.5.1
torch_npu	2.5.1
NumPy	1.26.4

环境配置

创建 Conda 环境

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

克隆代码

git clone https://atomgit.com/AI4Science/proteinbert_pytorch.git
cd proteinbert_pytorch
mkdir ./proteinbert_models

安装依赖

pip install torch==2.5.1 torch_npu==2.5.1 numpy==1.26.4 pyyaml \
    pandas scikit-learn h5py scipy
pip install decorator attrs psutil absl-py cloudpickle ml-dtypes tornado

如果不使用 Ascend NPU，只做 CPU 推理，则 torch_npu 不是必须的。

配置昇腾环境

source ./bin/env_npu.sh

该脚本会自动：

• source /usr/local/Ascend/ascend-toolkit/set_env.sh
• 补充 Ascend driver / CANN 的运行时库路径
• 默认设置 ASCEND_RT_VISIBLE_DEVICES=0

安装包

python setup.py install

模型权重准备

权重存放位置

./proteinbert_models/
├── epoch_92400_sample_23500000.pkl    # 原始 TF 预训练权重
└── proteinbert_pytorch.pt             # 转换后的 PyTorch checkpoint

下载原始 TensorFlow 权重

从以下地址下载：

• https://modelers.cn/models/Ascend-AI4S/proteinbert
• 或 GitHub: https://github.com/nadavbra/proteinbert_data_files/blob/master/epoch_92400_sample_23500000.pkl

转换为 PyTorch 权重

convert_tf_to_pytorch \
    --input ./proteinbert_models/epoch_92400_sample_23500000.pkl \
    --output ./proteinbert_models/proteinbert_pytorch.pt

推理命令

CPU 推理

inference_proteinbert_npu \
    --weights ./proteinbert_models/proteinbert_pytorch.pt \
    --seqs MKTVRQERLKSIVRILERSKEPVSGAQ ACDEFGHIKLMNPQRSTUVWXY \
    --device cpu

NPU 推理

inference_proteinbert_npu \
    --weights ./proteinbert_models/proteinbert_pytorch.pt \
    --seqs MKTVRQERLKSIVRILERSKEPVSGAQ ACDEFGHIKLMNPQRSTUVWXY \
    --device npu:0

快速代码示例

from proteinbert import convert_tf_to_pytorch, tokenize_seqs
import torch

seqs = [
    "MKTVRQERLKSIVRILERSKEPVSGAQ",
    "ACDEFGHIKLMNPQRSTUVWXY",
]
seq_len = 512

model, n_annotations = convert_tf_to_pytorch(
    "./proteinbert_models/epoch_92400_sample_23500000.pkl"
)
model = model.to("cpu").eval()

tokens = torch.from_numpy(tokenize_seqs(seqs, seq_len)).long()
input_annotations = torch.zeros(len(seqs), n_annotations)

with torch.no_grad():
    local_outputs, global_outputs = model(tokens, input_annotations)

Benchmark 微调

单个 Benchmark 微调

例如在 signalP_binary 上微调：

finetune_proteinbert_npu \
    --weights ./proteinbert_models/proteinbert_pytorch.pt \
    --train-csv ./protein_benchmarks/signalP_binary.train.csv \
    --test-csv ./protein_benchmarks/signalP_binary.test.csv \
    --task binary \
    --device npu:0

运行全部 Benchmark

HOME=$PWD TORCH_DEVICE_BACKEND_AUTOLOAD=0 python \
    ./scripts/demo_scripts/demo2_all_benchmarks_npu.py \
    --device npu:0

指定部分 benchmark：

HOME=$PWD TORCH_DEVICE_BACKEND_AUTOLOAD=0 python \
    ./scripts/demo_scripts/demo2_all_benchmarks_npu.py \
    --device npu:0 \
    --benchmarks signalP_binary fluorescence scop

支持的 Benchmark 列表

Benchmark	任务类型	指标
signalP_binary	全局二分类	AUC
fluorescence	全局回归	Spearman
remote_homology	全局多分类	Accuracy
stability	全局回归	Spearman
scop	全局多分类	Accuracy
secondary_structure	序列级分类	Accuracy
disorder_secondary_structure	序列级二分类	AUC
ProFET_NP_SP_Cleaved	全局二分类	AUC

精度验证

Embedding 提取验证

项目	结果
Local embedding 形状	(3, 512, 1562)
Global embedding 形状	(3, 15599)
CPU/NPU seq_probs 偏差	max_abs_diff=0.01190776, mean_abs_diff=0.00008779
CPU/NPU annotations 偏差	max_abs_diff=0.00603402, mean_abs_diff=0.00000099

NPU 推理性能

项目	结果
Batch size	4
Sequence length	512
平均推理时间	0.0048 s
吞吐量	~826 seq/s

下游 Benchmark 实测结果（NPU）

Benchmark	指标	NPU 实测结果
signalP_binary	AUC	0.995857
fluorescence	Spearman	0.662879
remote_homology	Accuracy	0.213092
stability	Spearman	0.765755
scop	Accuracy	0.885998
secondary_structure	Accuracy	0.740793
disorder_secondary_structure	AUC	0.872198
ProFET_NP_SP_Cleaved	AUC	0.982882

与 TensorFlow/GPU 基线对比

Benchmark	指标	GPU (TF)	NPU (PyTorch)	偏差
signalP_binary	AUC	0.9961	0.9965	+0.04%
fluorescence	Spearman	0.6475	0.6597	+1.22%
remote_homology	Accuracy	22.42%	21.17%	-1.25%
stability	Spearman	0.7068	0.7851	+7.83%
ProFET_NP_SP_Cleaved	AUC	0.9855	0.9852	-0.03%

分类任务整体与 TensorFlow 基线接近。回归任务的差异更大，主要受 TensorFlow Adam 与 PyTorch Adam 优化器收敛差异影响。

迁移适配要点

TensorFlow → PyTorch 关键差异

问题	解决方案
Keras 权重存储为 pkl dump	解析 pkl 文件提取各层 numpy 数组，按层名映射到 PyTorch
TF LayerNormalization 默认 eps=1e-3	PyTorch LayerNorm 中设置 eps=1e-3 保持一致
TF Conv1D 权重 (K, Cin, Cout)	转置为 PyTorch (Cout, Cin, K)
TF Dense 权重 (in, out)	转置为 PyTorch (out, in)
TF 全局注意力参数形状差异	按多头结构 reshape 后映射

昇腾 NPU 适配

PyTorch 版本通过 torch_npu 在昇腾 NPU 上运行，只需指定设备即可：

--device npu:0

CPU/NPU 精度偏差在可接受范围内（mean_abs_diff < 0.0001）。

已知限制

• 当前版本不包含原始仓库完整的 UniRef 预训练流水线的 PyTorch 复刻
• PhosphositePTM benchmark 因缺少 PhosphositePTM.train.csv 无法完成微调评测
• 回归任务（fluorescence、stability）因优化器差异，PyTorch 版与 TF 版结果偏差较大
• 建议与原始 TensorFlow 版使用不同 Conda 环境

参考文献

• Brandes, N., Ofer, D., Peleg, Y., Rappoport, N. & Linial, M. ProteinBERT: A universal deep-learning model of protein sequence and function. Bioinformatics (2022). https://doi.org/10.1093/bioinformatics/btac020
• 上游代码仓库：https://github.com/nadavbra/protein_bert
• PyTorch 昇腾适配版：https://atomgit.com/AI4Science/proteinbert_pytorch