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

ai-for-science-deepfri

Agents & models

View repository

Updated 10-15 17:19

About this skill

DeepFRI Ascend NPU Migration Skill

Project Overview

DeepFRI is a protein function prediction framework based on a GCN + LSTM language model, with the original implementation based on TensorFlow/Keras. This Skill documents the entire process of migrating it completely to PyTorch + Ascend NPU.

Model architecture:

  • DeepCNN: 16-way parallel Conv1D → BatchNorm → ReLU → GlobalMaxPool → FuncPredictor
  • DeepFRI GCN：LSTM LM → Embedding → 3×MultiGraphConv → SumPool → Dense → FuncPredictor
  • LSTM language model: 2-layer CuDNNLSTM (hidden=512), outputs concatenated into 1024-dimensional features

Supported Ontologies: MF (Molecular Function), BP (Biological Process), CC (Cellular Component), EC (Enzyme Classification)

Prerequisites

Item Requirement
Hardware Ascend910 series (at least 1 card)
CANN ≥ 8.2 (8.3.RC1 recommended)
Python 3.10
PyTorch 2.5.1
torch_npu 2.5.1

Migration Steps

Step 1: Clone repository & download model

cd /home/panjingxu
git clone https://github.com/flatironinstitute/DeepFRI.git
cd DeepFRI

# 下载 GPU 版预训练模型（README 精度参考值基于此版本）
wget https://users.flatironinstitute.org/~renfrew/DeepFRI_data/trained_models.tar.gz
tar xzf trained_models.tar.gz

Note: The official website provides two versions:

  • trained_models.tar.gz(GPU version): CNN + MultiGraphConv + CuDNNLSTM, the README accuracy reference values are based on this
  • newest_trained_models.tar.gz(CPU version): CNN + GraphConv + standard LSTM, different structure, different number of goterms

Step 2: Create Conda environment

export PIP_INDEX_URL=https://repo.huaweicloud.com/repository/pypi/simple/

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

# TF（仅用于权重提取，可选）
pip install tensorflow numpy==1.26.4 biopython scikit-learn

# PyTorch + NPU
pip install torch==2.5.1 torch_npu==2.5.1
pip install pyyaml decorator attrs psutil cloudpickle tornado

Step 3: Verify NPU availability

source /home/Ascend/ascend-toolkit/set_env.sh
python3 -c "import torch; import torch_npu; a = torch.randn(3,4).npu(); print(a+a)"

Step 4: Analyze TF model weight structure

Since newer versions of TF 2.x cannot load CuDNNLSTM, use h5py directly to read the HDF5 weights:

import h5py
f = h5py.File('trained_models/DeepFRI-MERGED_xxx.hdf5', 'r')
mw = f['model_weights']
for layer_name in mw.keys():
    grp = mw[layer_name]
    def print_ds(name, obj):
        if isinstance(obj, h5py.Dataset):
            print(f"  {name}: {obj.shape}")
    grp.visititems(print_ds)
f.close()

Step 5: Deploy PyTorch model files & convert weights

Copy the files under this Skill's scripts/ directory to the corresponding locations in the DeepFRI repository:

SKILL_DIR=/path/to/awesome-ascend-skills/skills/ai-for-science/models/deepfri/scripts
REPO_DIR=/path/to/DeepFRI   # 替换为实际路径

# 复制 PyTorch 模型文件到 deepfrier 包内
cp $SKILL_DIR/torch_layers.py    $REPO_DIR/deepfrier/
cp $SKILL_DIR/torch_model.py     $REPO_DIR/deepfrier/
cp $SKILL_DIR/torch_predictor.py $REPO_DIR/deepfrier/

# 复制脚本到仓库根目录
cp $SKILL_DIR/convert_weights.py $REPO_DIR/
cp $SKILL_DIR/predict_npu.py     $REPO_DIR/
cp $SKILL_DIR/verify_accuracy.py $REPO_DIR/

Then convert the weights:

cd $REPO_DIR
python convert_weights.py
# 生成 trained_models/pytorch/*.pt

scripts directory file list:

File Target location
scripts/torch_layers.py deepfrier/torch_layers.py
scripts/torch_model.py deepfrier/torch_model.py
scripts/torch_predictor.py deepfrier/torch_predictor.py
scripts/convert_weights.py convert_weights.py
scripts/predict_npu.py predict_npu.py
scripts/verify_accuracy.py verify_accuracy.py

Step 6: NPU inference

source /home/Ascend/ascend-toolkit/set_env.sh

# CNN 序列输入（对应 README Option 2）
python predict_npu.py \
  --seq 'SMTDLLSAEDIKKAIGAFTAADSFDHKKFFQMVGLKKKSADDVKKVFHILDKDKDGFIDEDELGSILKGFSSDARDLSAKETKTLMAAGDKDGDGKIGVEEFSTLVAES' \
  -ont mf --verbose --device npu:0

# CNN fasta 输入（对应 README Option 3）
python predict_npu.py \
  --fasta_fn examples/pdb_chains.fasta \
  -ont mf --verbose --device npu:0

Step 7: Accuracy verification

python verify_accuracy.py

Key conversion rules

Dense / Linear

TF PyTorch Convert
kernel: (in, out) weight: (out, in) Transpose .T
bias: (out,) bias: (out,) Copy directly

Conv1D

TF PyTorch Convert
Input (batch, length, channels) Input (batch, channels, length) transpose before and after
kernel: (K, Cin, Cout) weight: (Cout, Cin, K) np.transpose(w, (2,1,0))
bias: (Cout,) bias: (Cout,) Copy directly

BatchNorm (⚠️ the most critical pitfall in this migration)

TF PyTorch Convert
gamma weight Copy directly
beta bias Copy directly
moving_mean running_mean Copy directly
moving_variance running_var Copy directly
Default eps = 1e-3 Default eps = 1e-5 Must set eps=1e-3

⚠️ Fatal Bug: TF BatchNorm default eps=1e-3, PyTorch default eps=1e-5, 100 times difference. When moving_variance contains values close to 0, the wrong eps causes division-by-zero explosion, softmax saturates to 0/1,model output is all 0. This is the most critical discovery in this CNN migration.

CuDNNLSTM → nn.LSTM

TF CuDNNLSTM PyTorch nn.LSTM Convert
kernel: (input, 4*H) weight_ih: (4*H, input) Transpose .T
recurrent_kernel: (H, 4*H) weight_hh: (4*H, H) Transpose .T
bias: (8*H,) bias_ih: (4*H,) + bias_hh: (4*H,) Split into first half + second half 
pt_weight_ih = tf_kernel.T
pt_weight_hh = tf_recurrent_kernel.T
pt_bias_ih = tf_bias[:4 * H]   # 前半：input-hidden bias
pt_bias_hh = tf_bias[4 * H:]   # 后半：hidden-hidden bias

Accuracy comparison results

CNN model vs official website README (NPU Ascend910)

Test case GO term Official reference value NPU output diff Status
Option 2: seq→mf (1S3P-A) GO:0005509 calcium ion binding 0.99769 0.99769 5e-6 ✅ PASS
Option 3: fasta→mf (1S3P-A) GO:0005509 calcium ion binding 0.99769 0.99769 5e-6 ✅ PASS
Option 3: fasta→mf (2J9H-A) GO:0004364 glutathione transferase 0.46937 0.46936 9e-6 ✅ PASS
Option 3: fasta→mf (2J9H-A) GO:0016765 transferase activity 0.19910 0.19913 2.6e-5 ✅ PASS

CNN full Ontology validation

Ontology Status Top prediction
MF ✅ OK GO:0005509 score=0.99769 (calcium ion binding)
BP ✅ OK GO:0051179 score=0.14491 (localization)
CC ✅ OK GO:0005829 score=0.23145 (cytosol)
EC ✅ OK No prediction exceeds the threshold (this protein is not an enzyme, as expected)

GCN model description

The GPU version of the GCN model embeds a CuDNNLSTM language model. CuDNNLSTM uses GPU-specific weights. When run on a standard LSTM (PyTorch nn.LSTM), the step-by-step floating-point differences are exponentially amplified through the 3-layer MultiGraphConv graph convolution, causing the final output to saturate. If GCN functionality is needed, it is recommended to download the official newest_trained_models.tar.gz(CPU version) and use the standard LSTM + GraphConv.

List of newly added files

File Lines Description
deepfrier/torch_layers.py 199 PyTorch version of graph convolution layers (MultiGraphConv, GraphConv, SAGEConv, NoGraphConv, ChebConv, SumPooling, FuncPredictor)
deepfrier/torch_model.py 124 PyTorch version of models (LSTMLanguageModel, DeepFRIGCN, DeepCNN)
deepfrier/torch_predictor.py 178 PyTorch version of the inference predictor (PredictorPyTorch, compatible with the original TF version interface)
convert_weights.py 147 TF HDF5 → PyTorch state_dict weight conversion (LSTM/GCN/CNN)
predict_npu.py 46 NPU inference entry script (including transfer_to_npu automatic migration injection)
verify_accuracy.py 72 Accuracy comparison validation script (compares against the official README reference values)

predict_npu.py (NPU entry script)

Inject at the top of the entry script transfer_to_npu, automatically mapping torch.cuda.* to torch.npu.*：

import torch_npu
from torch_npu.contrib import transfer_to_npu
from deepfrier.torch_predictor import PredictorPyTorch

convert_weights.py core logic

  1. Use h5py to read all Datasets in the model_weights group of the TF HDF5 model
  2. CuDNNLSTM: transpose kernel/recurrent_kernel, split bias (8*H,) into ih + hh
  3. MultiGraphConv: directly copy kernel (both TF and PyTorch use (in, out) matrix multiplication)
  4. Dense: transpose kernel .T
  5. Conv1D：kernel np.transpose(w, (2,1,0))
  6. BatchNorm: directly copy gamma/beta/moving_mean/moving_var

Key implementation of torch_layers.py

The _normalize method of MultiGraphConv corresponds to the three-way normalization in the TF version:

def _normalize(self, A, eps=1e-6):
    # 去除自环
    A = A.clone()
    diag = torch.diagonal(A, dim1=-2, dim2=-1)
    A = A - torch.diag_embed(diag)
    A_hat = A + eye
    deg = A_hat.sum(dim=2)
    D_asymm = torch.diag_embed(1.0 / (eps + deg))
    D_symm = torch.diag_embed(1.0 / (eps + deg.sqrt()))
    # 返回三路：[原始A, 非对称归一化, 对称归一化]
    return [A, D_asymm @ A_hat, D_symm @ A_hat @ D_symm]

The three results are each multiplied by the feature matrix, then concatenated, and then passed through a linear transformation and ELU activation.

Key adaptations of DeepCNN in torch_model.py

# BatchNorm eps 必须匹配 TF 默认值
self.bn = nn.BatchNorm1d(total_filters, eps=1e-3)

# Conv1D 输入需要 transpose
def forward(self, seq):
    x = seq.transpose(1, 2)  # (batch, length, ch) → (batch, ch, length)
    conv_outs = [conv(x) for conv in self.conv_layers]
    x = torch.cat(conv_outs, dim=1)
    x = self.bn(x)
    x = F.relu(x)
    x = x.max(dim=2)[0]  # GlobalMaxPool
    return self.func_predictor(x)

Migration lessons learned

Pitfalls encountered (sorted by severity)

  1. BatchNorm eps mismatch(most critical): TF default 1e-3, PyTorch default 1e-5, causing CNN output to be all 0. Fix:nn.BatchNorm1d(dim, eps=1e-3)

  2. CuDNNLSTM platform limitation: the GPU version model embeds CuDNNLSTM, which cannot be loaded and validated with TF on aarch64 + NPU. Solution: use h5py to read weights directly + numpy to manually validate single-step

  3. TF 2.x is incompatible with CuDNNLSTM: TF 2.21's load_model does not recognize CuDNNLSTM. Solution: bypass model loading and extract weights directly from HDF5

  4. Conv1D dimension order：TF (batch, length, channels) vs PyTorch (batch, channels, length), transpose is required before and after

General lessons applicable to other projects

  • When a TF model cannot be loaded, prefer to use h5py to directly read the HDF5 weight structure
  • All normalization layers (LayerNorm, BatchNorm) must have their default eps values checked
  • CuDNNLSTM bias (8*H,) Splitting rule: the first half = input bias, the second half = recurrent bias
  • After conversion, first compare the intermediate layer outputs (mean/std/min/max) to locate the first diverging layer

Known limitations

  • The GCN model (cmap input) uses a CuDNNLSTM language model; these weights are GPU-specific, and a standard LSTM will exhibit numerical amplification in deep MultiGraphConv. If GCN functionality is needed, it is recommended to use the official CPU version of the model
  • CNN padding='same' PyTorch will issue a warning for even kernel sizes; this does not affect accuracy
  • Ascend910 does not support fp64; torch_npu automatically downgrades to fp32, which does not affect this model

References