ProteinBERT模型在昇腾平台的迁移pytorch实践

weixin_45870772

76人浏览 · 2026-07-31 11:13:18

weixin_45870772 · 2026-07-31 11:13:18 发布

作者：昇腾实战派

知识地图：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	捕获蛋白质整体功能特征

模型组成

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

每个 Block 包含：
- 局部路径：窄卷积（kernel=9, dilation=1）+ 宽卷积（kernel=9, dilation=5）+ 位置全连接层 + LayerNorm
- 全局路径：两层全连接 + LayerNorm
- 局部→全局：全局注意力层（Global Attention），全局表示作为 Query，局部表示作为 Key/Value
- 全局→局部：广播全连接层，将全局信息投影到每个序列位置
输出层
- 局部输出：序列每个位置的 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）

微调协议：

冻结所有预训练层，仅训练新增全连接层 40 epochs
解冻所有层，继续训练 40 epochs
在更大序列长度上做最终 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

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

更多推荐

安全底座：企业服务赛道真正的隐性分水岭

在供应链安全与信创适配方面，BeeWorks 已完成与国产主流基础软硬件的全面适配，覆盖飞腾、鲲鹏、龙芯等国产芯片，银河麒麟、统信UOS等国产操作系统，以及达梦、人大金仓等国产数据库，能够为企业构建从底层硬件到上层应用的全国产化安全底座。单一环节的安全强化无法形成有效防护，只有将身份管理、访问控制、通信加密、操作审计四大能力整合为统一的安全门户，才能真正建立系统性的安全防线。在安全审查成为刚性门槛