AI For Science Deepfri Tf Npu

DeepFRI TensorFlow 原生昇腾 NPU 迁移 Skill

项目概述

DeepFRI 是 Flatiron Institute 开发的蛋白质功能预测框架，原始实现基于 TensorFlow/Keras。本 Skill 记录使用 TF Community 方式（TF 2.6.5 + npu_device） 将其直接部署到昇腾 NPU 上，无需转换为 PyTorch。

与 TF→PyTorch 迁移方式的区别：

模型架构：

• DeepCNN：序列输入 → 多路 Conv1D → BatchNorm → ReLU → GlobalMaxPool → FuncPredictor
• DeepFRI GCN：接触图+序列 → LSTM LM Embedding → MultiGraphConv → SumPool → FuncPredictor

支持的 Ontology：MF（分子功能）、BP（生物过程）、CC（细胞组分）、EC（酶分类）

前置条件

迁移步骤

Step 1：克隆仓库 & 下载模型

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

模型版本选择：

• newest_trained_models.tar.gz（CPU 版，推荐）：标准 LSTM + GraphConv，兼容 NPU
• trained_models.tar.gz（GPU 版）：含 CuDNNLSTM，仅 CNN 部分可正常使用

Step 2：搭建 TF 2.6.5 + npu_device 环境

此步骤较复杂，详见 ascend-tf-community Skill 的 Step 1-3。 以下为快速摘要：

# 创建环境
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：适配代码

DeepFRI 使用标准 TF 2.x Keras API，适配量极小。只需修改 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")

无需修改的文件（直接兼容 NPU）：

• deepfrier/layers.py — 自定义图卷积层
• deepfrier/utils.py — 数据处理工具
• deepfrier/DeepFRI.py — 模型定义
• deepfrier/DeepCNN.py — CNN 模型定义

Step 4：运行 NPU 推理

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：精度验证

创建 verify_accuracy.py，先 CPU 跑一遍得到基线，再 NPU 跑一遍对比：

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

SEQ = "MKFLILLFNILCLFPVLAADNHGVSMQTESGHLVHEVTLHALTDKDLFGKVRANALTK" "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")

精度对比结果

CNN 模型（-s 序列输入，CPU vs NPU Ascend910）

说明：

• 预测排名（GO term 顺序）完全一致
• 分数差异在 NPU HF32 混合精度的正常范围内
• 如需更高精度，设置 npu_device.global_options().precision_mode = "allow_fp32"

NPU 推理输出示例（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 模型（--cmap / -pdb 接触图输入，CPU vs NPU Ascend910）

GCN 模型需要 CuDNNLSTM → 标准 LSTM 权重转换（自动处理，详见下文）。

关键适配点

1. npu_device 初始化

import npu_device
npu_device.open().as_default()

• 必须在 import tensorflow 之后、模型加载之前调用
• 调用后 TF 默认设备变为 NPU，所有 Keras 算子自动调度到 NPU

2. CuDNNLSTM → 标准 LSTM 自动转换（GCN 模型核心修改）

GPU 版 GCN 模型内嵌 CuDNNLSTM 语言模型，该算子在 NPU/CPU 上均无 kernel。 修改 Predictor.py 添加自动兼容加载逻辑：

1. 检测 HDF5 中 LSTM bias 形状是否为 (8*H,)（CuDNNLSTM 特征）
2. Monkey-patch CuDNNLSTM 类为标准 tf.keras.layers.LSTM wrapper
3. 从 config 重建模型架构（不加载权重）
4. 用 load_weights(by_name=True, skip_mismatch=True) 加载非 LSTM 权重
5. 手动转换 LSTM 权重：bias[:4*H] + bias[4*H:] 合并为 (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. 自定义 Layer 兼容性

DeepFRI 的自定义层（MultiGraphConv、GraphConv、FuncPredictor、SumPooling） 均使用标准 TF 原语（tf.matmul、tf.reduce_sum、tf.keras.backend.batch_dot 等）， npu_device 可自动处理，无需任何修改。

3. 模型加载方式

DeepFRI 使用 tf.keras.models.load_model 加载 HDF5 模型，通过 custom_objects 注册自定义层。此方式完全兼容 npu_device，模型权重自动放置到 NPU。

4. 推理模式

DeepFRI 使用 model([inputs], training=False) 进行推理。training=False 确保 BatchNorm 使用 moving_mean/moving_variance 而非批统计量，NPU 上行为一致。

迁移经验总结

本次迁移的关键发现

1. TF Community 方式代码改动极小：整个 DeepFRI 项目仅需修改 predict.py 添加 3 行代码（import + open + 参数），所有其他文件零改动

2. 自定义 Keras Layer 自动兼容：npu_device 通过 TF 的 op 调度机制接管计算， 用户定义的 tf.keras.layers.Layer 子类无需任何适配

3. HF32 精度影响可控：Ascend910 默认使用 HF32 混合精度，对 DeepFRI 的 预测排名无影响，仅绝对分数有 ~1e-2 级别差异

4. 编译 TF 是主要工作量：aarch64 上编译 TF 2.6.5 约 4-5 分钟，但环境配置 （Bazel 版本、nsync 补丁、ABI 设置、GitHub 代理）需要仔细处理

与 TF→PyTorch 迁移方式的经验对比

文件变更清单

已知限制

• aarch64 必须源码编译 TF 2.6.5（x86 可直接 pip install）
• GPU 版预训练模型中的 CuDNNLSTM 在 TF 2.6 上无法加载，需使用 CPU 版模型
• npu_device 的 Unrecognized graph engine option 警告可忽略（CANN 版本差异）
• 首次 NPU 推理会有算子编译开销（~10-20 秒），后续推理正常速度
• 动态 shape 算子（如变长序列）部分可能 fallback 到 CPU

依赖的 Skill

• ascend-tf-community：通用昇腾 TF Community 迁移手册（TF 编译、tfplugin 安装流程）

参考资料

• DeepFRI TensorFlow 原生运行检查表：references/tf-runtime-checklist.md