SuperAPI AI Directory
Partner with us中文AI newsFavorites
Skillsai-for-science-proteinbert

ai-for-science-proteinbert

Agents & models

View repository

Updated 10-15 17:19

About this skill

ProteinBERT Ascend NPU Deployment Skill

Fully migrate ProteinBERT from TensorFlow/Keras to PyTorch + torch_npu, accuracy has been verified on 5 benchmark tasks (with GPU as the baseline).

Model Overview

Item Value
Architecture 6 layers Conv + GlobalAttention blocks
Parameter count ~16M (vocab=26, annotations=8943, d_seq=128, d_global=512)
Original framework TensorFlow 2.x / Keras
Target framework PyTorch 2.5.1 + torch_npu 2.5.1
Verified CANN 8.3.RC1
Hardware Ascend910

Environment Preparation

# CANN 环境
source /home/Ascend/ascend-toolkit/set_env.sh
export ASCEND_RT_VISIBLE_DEVICES=0

# Conda 环境
conda create -n proteinbert_npu python=3.11 -y
conda activate proteinbert_npu
pip install torch==2.5.1 torch_npu==2.5.1 numpy==1.26.4 pyyaml \
    pandas scikit-learn h5py scipy -i https://repo.huaweicloud.com/repository/pypi/simple/

File Structure

This Skill is self-contained with all code and can be copied and used directly:

proteinbert/
├── SKILL.md                            # 本文件
├── references/
│   └── migration_details.md              # 详细迁移参考（权重映射、bug 记录、调试方法）
└── scripts/
    ├── proteinbert_pytorch/              # PyTorch 核心实现
    │   ├── model.py                       #   模型定义
    │   ├── convert_weights.py             #   TF pkl → PyTorch 权重转换
    │   ├── inference.py                   #   分词器 + 推理工具
    │   └── finetune.py                    #   微调框架
    ├── demo_scripts/                      # GPU/NPU 配对用例脚本
    │   ├── demo1_signalP_gpu.py           #   用例1: signalP 微调 (TF)
    │   ├── demo1_signalP_npu.py           #   用例1: signalP 微调 (PyTorch)
    │   ├── demo2_all_benchmarks_gpu.py    #   用例2: 5个benchmark (TF)
    │   ├── demo2_all_benchmarks_npu.py    #   用例2: 5个benchmark (PyTorch)
    │   ├── demo3_attention_gpu.py         #   用例3: 注意力可视化 (TF)
    │   └── demo3_attention_npu.py         #   用例3: 注意力可视化 (PyTorch)
    ├── tools/                             # 精度对比工具
    │   ├── get_embeddings_gpu.py          #   Embedding提取 (TF)
    │   ├── get_embeddings_npu.py          #   Embedding提取 (PyTorch)
    │   ├── debug_layerwise_gpu.py         #   逐层调试 (TF)
    │   ├── debug_layerwise_npu.py         #   逐层调试 (PyTorch)
    │   └── compare_embeddings.py          #   GPU/NPU 结果对比
    └── deploy_toolkit/                    # 独立部署包
        ├── convert_weights.py             #   权重转换（自包含模型定义）
        ├── inference_npu.py               #   NPU 推理
        ├── finetune_npu.py                #   NPU 微调
        ├── setup.sh                       #   一键部署脚本
        └── requirements.txt               #   依赖列表

Quick Start

Step 1: Download the pretrained weights

Download from Zenodo (183MB):

https://zenodo.org/records/10371965/files/epoch_92400_sample_23500000.pkl

Place into ~/proteinbert_models/epoch_92400_sample_23500000.pkl。

Step 2: Convert the weights

python scripts/deploy_toolkit/convert_weights.py \
    --input ~/proteinbert_models/epoch_92400_sample_23500000.pkl \
    --output ~/proteinbert_models/proteinbert_pytorch.pt

Step 3: Extract Embedding

Original GPU version (TF):

from proteinbert import load_pretrained_model
from proteinbert.conv_and_global_attention_model import get_model_with_hidden_layers_as_outputs

pretrained_model_generator, input_encoder = load_pretrained_model()
model = get_model_with_hidden_layers_as_outputs(pretrained_model_generator.create_model(seq_len))
encoded_x = input_encoder.encode_X(seqs, seq_len)
local_representations, global_representations = model.predict(encoded_x, batch_size=batch_size)

NPU equivalent version (PyTorch):

from proteinbert_pytorch.convert_weights import convert_tf_to_pytorch
from proteinbert_pytorch.inference import tokenize_seqs
import torch, torch_npu
from torch_npu.contrib import transfer_to_npu

model, n_ann = convert_tf_to_pytorch('~/proteinbert_models/epoch_92400_sample_23500000.pkl')
model = model.to('npu:0').eval()
tokens = tokenize_seqs(seqs, seq_len)
inp_seq = torch.from_numpy(tokens).long().to('npu:0')
inp_ann = torch.zeros(len(seqs), n_ann).to('npu:0')
with torch.no_grad():
    output_seq, output_ann = model(inp_seq, inp_ann)

Full hidden layer concatenation version (benchmarked against get_model_with_hidden_layers_as_outputs) see scripts/tools/get_embeddings_npu.py. Output dimensions:

  • local_representations: (batch, seq_len, 1562) = 12 LayerNorm×128 + output_seq×26
  • global_representations: (batch, 15599) = input_dense×512 + 12 LayerNorm×512 + output_ann×8943

Step 4: Fine-tuning training

All three use cases provide paired GPU/NPU scripts, with consistent print formats for easy diff comparison.

Use case 1: signalP_binary fine-tuning

python scripts/demo_scripts/demo1_signalP_gpu.py   # GPU端
python scripts/demo_scripts/demo1_signalP_npu.py   # NPU端

Three-stage training: frozen layers (lr=1e-2) → all layers (lr=1e-4) → long sequences (seq_len=1024, lr=1e-5)

Use case 2: Run all 5 Benchmarks

python scripts/demo_scripts/demo2_all_benchmarks_gpu.py   # GPU端
python scripts/demo_scripts/demo2_all_benchmarks_npu.py   # NPU端

Includes signalP_binary, fluorescence, remote_homology, stability, ProFET_NP_SP_Cleaved.

Use case 3: Attention visualization

python scripts/demo_scripts/demo3_attention_gpu.py   # GPU端
python scripts/demo_scripts/demo3_attention_npu.py   # NPU端

Extract the weights of 24 attention heads (6 layers×4 heads), save as npz for easy numerical comparison.

Accuracy Verification Results

Benchmark fine-tuning accuracy (with GPU/TF as the baseline)

Benchmark Task type Metric GPU (TF) NPU (PyTorch) Deviation Status
signalP_binary Binary classification AUC 0.9961 0.9965 +0.04% ✅
fluorescence Regression Spearman 0.6475 0.6597 +1.22% ✅
remote_homology Multiclass classification Accuracy 22.42% 21.17% -1.25% ✅
stability Regression Spearman 0.7068 0.7851 +7.83% ✅
ProFET_NP_SP_Cleaved Binary classification AUC 0.9855 0.9852 -0.03% ✅

The deviations for binary/multiclass classification tasks are all within 1.3%. For regression tasks, the NPU deviation is larger but in the positive direction (better), the reason is the optimizer convergence difference between PyTorch Adam and TF Adam.

Layer-by-layer comparison of pretrained model Embedding (no training randomness)

Layer name GPU mean NPU mean Deviation
embedding-seq-input 0.000829 0.000829 0
dense-global-input -0.005728 -0.005728 0
global-merge1-norm-block1 -0.031465 -0.031465 <1e-5
global-merge2-norm-block6 -0.010764 -0.010764 <1e-5
output-annotations 0.000148 0.000148 <1e-5

The mean/std deviations of the outputs of all 28 layers are < 1e-4, and the weight conversion values are completely consistent.

Key Migration Adaptation Points

For detailed weight mapping, code comparison, and debugging methods, see references/migration_details.md。

Critical level (causes completely incorrect accuracy before fixing)

  1. K.dot and einsum dimension mapping error

    • TF K.dot(A, B) contracting the last dimension of A and the second-to-last dimension of B
    • einsum('bls,hvs->...') is wrong, the correct way to write it is einsum('bls,hsv->...')
    • Reversing it causes the output to be completely wrong starting from the 1st block

  2. LayerNorm epsilon default value differs by 100 times

    • TF default eps=1e-3, PyTorch default eps=1e-5
    • must be set explicitly nn.LayerNorm(dim, eps=1e-3)

Critical (causes shape or training behavior anomalies)

  1. Hidden layer collection logic mismatch

    • TF filter name input-seq-encoding and actual layer name embedding-seq-input do not match
    • causing the seq dimension to be 1562 instead of 1690

  2. Incorrect input dimension for the fine-tuning head

    • The TF version head receives the concatenation of all layers at 15599 dimensions, rather than the last layer's 512 dimensions
    • Using the wrong dimension causes training to trigger Early Stopping prematurely

Normal (expected differences between frameworks)

  1. pickle weight order：global_input_dense in embedding before
  2. Conv1D weights：TF (kernel,in,out) → PyTorch (out,in,kernel) need transpose
  3. Dense weights：TF (in,out) → PyTorch (out,in) need transpose
  4. pyyaml dependency: required by the torch_npu runtime, not listed in the official dependencies
  5. TF/torch_npu conflict: cannot be installed simultaneously in the same conda environment

Precision comparison method

# 步骤1: 分别在 GPU 和 NPU 上运行配对脚本
python scripts/tools/get_embeddings_gpu.py   # → embeddings_gpu.npz
python scripts/tools/get_embeddings_npu.py   # → embeddings_npu.npz

# 步骤2: 数值对比
python scripts/tools/compare_embeddings.py

# 步骤3: 如果偏差较大，用逐层调试定位
python scripts/tools/debug_layerwise_gpu.py > gpu.txt
python scripts/tools/debug_layerwise_npu.py > npu.txt
diff gpu.txt npu.txt

The first divergence point in the layer-by-layer comparison is the layer where the bug is located.

References