ai-for-science-deepfri-tf-npu

DeepFRI TensorFlow Native Ascend NPU Migration Skill

Project Overview

DeepFRI is a protein function prediction framework developed by the Flatiron Institute, with the original implementation based on TensorFlow/Keras. This Skill documents using the TF Community approach (TF 2.6.5 + npu_device) to deploy it directly on Ascend NPU without converting to PyTorch.

Differences from the TF-to-PyTorch migration approach:

Dimension	TF-to-PyTorch Migration	TF Community Migration (this Skill)
Amount of code changes	Large (rewrite model + convert weights)	Minimal (only add 2 lines of initialization)
Dependent framework	PyTorch + torch_npu	TensorFlow 2.6.5 + npu_device
Compilation requirement	None	Requires source code compilation of TF (aarch64)
GCN model support	Requires additional handling of CuDNNLSTM	Directly supported (CPU version model)
Precision	< 1e-4 (FP32 exact match)	< 1e-2 (HF32 mixed precision)

Model architecture:

• DeepCNN: sequence input → multi-branch Conv1D → BatchNorm → ReLU → GlobalMaxPool → FuncPredictor
• DeepFRI GCN: contact map + sequence → LSTM LM Embedding → MultiGraphConv → SumPool → FuncPredictor

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.0 (tested with 8.5 working)
Python	3.8.x
TensorFlow	2.6.5 (aarch64 requires source code compilation)
npu_device	0.1 (from Ascend-cann-tfplugin 8.0.RC3)
Bazel	3.7.2

---

Migration Steps

Step 1: Clone repository & download models

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

# 下载 CPU 版预训练模型（推荐，使用标准 LSTM，兼容性好）
wget https://users.flatironinstitute.org/~renfrew/DeepFRI_data/newest_trained_models.tar.gz
tar xzf newest_trained_models.tar.gz

Model version selection：

• newest_trained_models.tar.gz(CPU version, recommended): standard LSTM + GraphConv, compatible with NPU
• trained_models.tar.gz(GPU version): contains CuDNNLSTM, only the CNN part can be used normally

Step 2: Set up the TF 2.6.5 + npu_device environment

This step is relatively complex, see ascend-tf-community Skill's Step 1-3. The following is a quick summary:

# 创建环境
conda create -n deepfri_npu python=3.8 -y
conda activate deepfri_npu
conda install -y hdf5

# 安装 Bazel 3.7.2
wget 'https://mirrors.huaweicloud.com/bazel/3.7.2/bazel-3.7.2-linux-arm64' \
     -O /usr/local/bin/bazel && chmod +x /usr/local/bin/bazel

# 编译安装 TF 2.6.5（aarch64, ABI=0, 含 nsync 补丁）
# 详见 ascend-tf-community Skill Step 2

# 安装 tfplugin npu_device
# 详见 ascend-tf-community Skill Step 3

# 安装 DeepFRI 依赖
pip install numpy==1.21.6 biopython scikit-learn networkx h5py==3.7.0
pip install decorator sympy scipy attrs psutil

Step 3: Adapt the code

DeepFRI uses the standard TF 2.x Keras API, so the amount of adaptation is minimal. Only modify predict.py：

# predict.py 头部添加 --npu 参数支持
import os
import npu_device  # 新增

# 在 argparse 中添加：
parser.add_argument('--npu', action='store_true', help="Run on Ascend NPU.")

# 在 args = parser.parse_args() 之后添加：
if args.npu:
    npu_device.open().as_default()
    print("### Running on Ascend NPU")

Files that do not need modification(directly compatible with NPU):

• deepfrier/layers.py — custom graph convolution layer
• deepfrier/utils.py — data processing utilities
• deepfrier/DeepFRI.py — model definition
• deepfrier/DeepCNN.py — CNN model definition

Step 4: Run NPU inference

source /home/Ascend/8.5/cann-8.5.0/set_env.sh

# CNN 序列推理（MF）
python predict.py --npu \
    -s 'MKFLILLFNILCLFPVLAADNHGVSMQTESGHLVHEVTLHALTDKDLFGKVRANALTK' \
    -ont mf -o npu_result -v

# CNN 序列推理（全部 Ontology）
python predict.py --npu \
    -s 'MKFLILLFNILCLFPVLAADNHGVSMQTESGHLVHEVTLHALTDKDLFGKVRANALTK' \
    -ont mf bp cc ec -o npu_result -v

# GCN PDB 文件推理（需要 CPU 版模型）
python predict.py --npu \
    -pdb examples/1S3P-A.pdb \
    -ont mf -o npu_gcn_result -v

Step 5: Precision verification

Create verify_accuracy.py, first run once on CPU to get a baseline, then run once on NPU for comparison:

#!/usr/bin/env python3
import os, json, numpy as np

SEQ = "MKFLILLFNILCLFPVLAADNHGVSMQTESGHLVHEVTLHALTDKDLFGKVRANALTK"
ONTOLOGIES = ["mf", "ec"]

def run_predict(use_npu, prefix):
    args = f"python predict.py -s '{SEQ}' -ont {' '.join(ONTOLOGIES)} -o {prefix} -v"
    if use_npu:
        args = args.replace("predict.py", "predict.py --npu")
    os.system(args)

def compare(cpu_prefix, npu_prefix):
    for ont in ONTOLOGIES:
        cpu_file = f"{cpu_prefix}_{ont.upper()}_pred_scores.json"
        npu_file = f"{npu_prefix}_{ont.upper()}_pred_scores.json"
        with open(cpu_file) as f: cpu_data = json.load(f)
        with open(npu_file) as f: npu_data = json.load(f)
        cpu_s = np.array(cpu_data["Y_hat"])
        npu_s = np.array(npu_data["Y_hat"])
        print(f"[{ont.upper()}] max_diff={np.max(np.abs(cpu_s-npu_s)):.6e} "
              f"cosine_sim={np.dot(cpu_s.flat,npu_s.flat)/(np.linalg.norm(cpu_s)*np.linalg.norm(npu_s)+1e-12):.8f}")

run_predict(False, "cpu_result")
run_predict(True, "npu_result")
compare("cpu_result", "npu_result")

---

Precision comparison results

CNN model (-s sequence input, CPU vs NPU Ascend910)

Ontology	max_diff	mean_diff	cosine_sim	Status
MF (molecular function)	2.37e-02	6.11e-04	0.99960	✅ PASS
EC (enzyme classification number)	3.83e-03	1.31e-04	0.99993	✅ PASS

Notes：

• Prediction ranking (GO term order) is completely identical
• The score difference is within the normal range for NPU HF32 mixed precision.
• For higher precision, set npu_device.global_options().precision_mode = "allow_fp32"

NPU inference output example (MF Ontology)

query_prot GO:0005215 0.25544 transporter activity
query_prot GO:0140318 0.23496 protein transporter activity
query_prot GO:0022857 0.22256 transmembrane transporter activity
query_prot GO:0015318 0.20254 inorganic molecular entity transmembrane transporter activity
query_prot GO:0008324 0.19642 cation transmembrane transporter activity

GCN model (--cmap / -pdb contact map input, CPU vs NPU Ascend910)

Protein	Input	Top-1 GO term (CPU)	Score CPU	Score NPU	Rank consistent
1S3P-A	npz cmap	GO:0032553 ribonucleotide binding	0.89507	0.88551	✅
1S3P-A	pdb file	GO:0032553 ribonucleotide binding	0.89507	0.88579	✅

The GCN model requires CuDNNLSTM → standard LSTM weight conversion (handled automatically, see below).

---

Key adaptation points

1. npu_device initialization

import npu_device
npu_device.open().as_default()

• Must be called after import tensorflow and before model loading
• After the call, the TF default device becomes NPU, and all Keras operators are automatically scheduled to the NPU.

2. CuDNNLSTM → standard LSTM automatic conversion (core modification for the GCN model)

The GPU version of the GCN model embeds a CuDNNLSTM language model, and this operator has no kernel on either NPU or CPU. Modify Predictor.py to add automatic compatibility loading logic:

1. Detect whether the LSTM bias shape in HDF5 is (8*H,)(CuDNNLSTM feature)
2. Monkey-patch CuDNNLSTM class to standard tf.keras.layers.LSTM wrapper
3. Rebuild the model architecture from config (without loading weights)
4. Use load_weights(by_name=True, skip_mismatch=True) to load non-LSTM weights
5. Manually convert LSTM weights:bias[:4*H] + bias[4*H:] merged into (4*H,) bias

def _convert_cudnn_weights_to_lstm(kernel, recurrent, bias):
    h = bias.shape[0] // 8
    return [kernel, recurrent, bias[:4*h] + bias[4*h:]]

3. Custom Layer compatibility

DeepFRI's custom layers (MultiGraphConv、GraphConv、FuncPredictor、SumPooling) all use standard TF primitives (tf.matmul、tf.reduce_sum、tf.keras.backend.batch_dot , etc.), npu_device can handle them automatically without any modification.

3. Model loading method

DeepFRI uses tf.keras.models.load_model to load the HDF5 model, and registers custom layers through custom_objects This method is fully compatible with npu_device, and model weights are automatically placed on the NPU.

4. Inference mode

DeepFRI uses model([inputs], training=False) for inference.training=False ensures BatchNorm uses moving_mean/moving_variance rather than batch statistics, and behavior on the NPU is consistent.

---

Migration experience summary

Key findings of this migration

1. The TF Community approach requires very few code changes: the entire DeepFRI project only needs modification of predict.py adding 3 lines of code (import + open + parameters), with zero changes to all other files

2. Custom Keras Layers are automatically compatible: npu_device takes over computation through TF's op scheduling mechanism, user-defined tf.keras.layers.Layer subclasses require no adaptation at all

3. The impact of HF32 precision is controllable: Ascend910 uses HF32 mixed precision by default, which has no impact on DeepFRI's prediction ranking, with only ~1e-2 level differences in absolute scores

4. Compiling TF is the main workload: compiling TF 2.6.5 on aarch64 takes about 4-5 minutes, but environment configuration (Bazel version, nsync patch, ABI settings, GitHub proxy) needs to be handled carefully

Experience comparison with the TF→PyTorch migration approach

Dimension	TF→PyTorch	TF Community
Workload	3-5 days (model rewrite + weight conversion + debugging)	0.5-1 day (compile TF + add 2 lines of code)
Precision risk	High (requires layer-by-layer alignment, many pitfalls such as BatchNorm eps)	Low (same TF code, only different compute hardware)
Long-term maintenance	Two codebases need to be maintained	Fully synchronized with the original TF code
Performance optimization headroom	Large (can be custom-optimized with torch_npu)	Limited by npu_device automatic optimization
Recommended scenarios	Projects that require the PyTorch ecosystem	Pure inference deployment, quick validation

---

File change list

File	Change	Description
predict.py	Modified	Added --npu parameter and npu_device.open() initialization (+6 lines)
deepfrier/Predictor.py	Modified	Added CuDNNLSTM→LSTM automatic compatible loading logic (+80 lines)
verify_accuracy.py	Added	CPU vs NPU accuracy comparison validation script

---

Known limitations

• aarch64 must compile TF 2.6.5 from source (x86 can pip install directly)
• The CuDNNLSTM in GPU-version pretrained models cannot be loaded on TF 2.6; the CPU-version model must be used
• npu_device's Unrecognized graph engine option warning can be ignored (CANN version difference)
• The first NPU inference incurs operator compilation overhead (~10-20 seconds); subsequent inference runs at normal speed
• Some dynamic shape operators (such as variable-length sequences) may fall back to CPU

---

Dependent Skills

• ascend-tf-community: General Ascend TF Community migration manual (TF compilation, tfplugin installation process)

Reference materials

• DeepFRI TensorFlow native run checklist:references/tf-runtime-checklist.md