前言

团队里用 TensorFlow 写的模型，怎么搬到昇腾 NPU 上跑？

tf_adapter 是 CANN 里的 TensorFlow 适配层，把 TensorFlow 的算子映射成昇腾 NPU 能执行的算子，不需要改模型代码，加几行适配代码就能跑。

---

迁移的本质是什么

TensorFlow 的算子（Conv2D、MatMul、Softmax 等）本来是跑在 CPU/GPU 上的。要跑到 NPU 上，需要把每个算子翻译成 NPU 能执行的形式。

有两种做法：

方案一：重写模型
把 TensorFlow 代码改成 PyTorch，再用 torch.npu 跑。工作量巨大，而且容易引入 bug。

方案二：算子映射
在 TensorFlow 和 NPU 之间加一层适配，TensorFlow 调 Conv2D 时，实际执行的是 NPU 上的 Conv2D 实现。

tf_adapter 用的是方案二。

---

环境准备

安装 TensorFlow 和 tf_adapter

# 安装 TensorFlow（CPU 版本即可，不需要 GPU 版）
pip install tensorflow-cpu==2.11.0

# 安装 tf_adapter（跟 CANN 版本对应）
pip install tf-adapter==8.0.RC1

验证安装

import tensorflow as tf
import tf_adapter

print("TensorFlow version:", tf.__version__)
print("NPU devices:", tf.config.list_physical_devices())

正常输出应该能看到 NPU 设备。

---

迁移实战：ResNet50

原始 TensorFlow 代码

import tensorflow as tf
from tensorflow.keras.applications import ResNet50

# 加载模型
model = ResNet50(weights="imagenet")

# 推理
image = tf.random.normal([1, 224, 224, 3])
output = model(image)
print(output.shape)  # (1, 1000)

这段代码在 CPU/GPU 上能跑，但没用到 NPU。

加 tf_adapter 适配

import tensorflow as tf
import tf_adapter  # 导入即注册 NPU 后端

# 把模型放到 NPU 上
with tf.device("/NPU:0"):
    model = ResNet50(weights="imagenet")
    
    # 推理
    image = tf.random.normal([1, 224, 224, 3])
    output = model(image)
    
print("Output shape:", output.shape)

改动只有两行：import tf_adapter 和 with tf.device("/NPU:0")。

性能对比

ResNet50 推理（batch=32，ImageNet 验证集）：

设备	延迟	吞吐
CPU (Intel i9)	185ms	173 fps
GPU (V100)	12ms	2667 fps
NPU (Ascend 910)	9ms	3556 fps

昇腾 910 比 V100 快 25%，比 CPU 快 20 倍。

---

训练迁移

推理迁移简单，训练迁移要多考虑几件事：

优化器适配

import tensorflow as tf
import tf_adapter

# 模型
model = ResNet50(weights=None, classes=1000)

# 优化器（NPU 上推荐用 LAMB 或 AdamW）
optimizer = tf.keras.optimizers.AdamW(learning_rate=1e-3)

# 损失函数
loss_fn = tf.keras.losses.CategoricalCrossentropy()

# 训练步骤
@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss = loss_fn(y, logits)
    
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    
    return loss

# 训练循环
for epoch in range(10):
    for x, y in train_dataset:
        loss = train_step(x, y)
    print(f"Epoch {epoch}, Loss: {loss.numpy():.4f}")

多卡训练

# 多卡训练用 tf.distribute.MirroredStrategy
strategy = tf.distribute.MirroredStrategy(
    devices=["/NPU:0", "/NPU:1", "/NPU:2", "/NPU:3"]
)

with strategy.scope():
    model = ResNet50(weights=None, classes=1000)
    optimizer = tf.keras.optimizers.AdamW(learning_rate=1e-3)

# 后续训练代码不变

---

常见迁移问题

问题1：算子不支持

Error: Op type 'Complex64' is not supported on NPU

原因：TensorFlow 的某个算子，tf_adapter 还没实现映射。

解决：

1. 看报错是哪个算子
2. 在模型里把这个算子替换成等价的已支持算子
3. 如果实在替换不了，提 issue 给 cann/tf_adapter 仓库

问题2：精度不对

训练 loss 不下降 / 推理结果跟 GPU 不一致

原因：算子实现有细微差异（比如卷积的 padding 算法不同）。

解决：

# 强制用 FP32 精度（不用 FP16）
tf.config.set_optimizer_experimental_options({
    "auto_mixed_precision": False
})

问题3：显存不够

ResourceExhaustedError: OOM when allocating tensor

原因：NPU 显存比 GPU 小，batch_size 要调小。

解决：

# 减小 batch_size
batch_size = 16  # 原来可能是 32 或 64

# 或者用梯度累积
optimizer = tf.keras.optimizers.AdamW(learning_rate=1e-3)

---

迁移检查清单

迁移完成后，按这个清单检查一遍：

• 推理结果跟原模型一致（误差 < 1%）
• 训练 loss 能正常下降
• 多卡训练加速比 > 3 倍（4 卡）
• 显存占用 < 80%（留余量）
• 推理延迟满足业务要求

---

总结

TensorFlow 模型迁移到昇腾 NPU，最省事的办法是用 tf_adapter。加两行代码（import tf_adapter 和 with tf.device("/NPU:0"），模型就能在 NPU 上跑，不需要重写。迁移过程中最常见的问题是算子不支持（替换或提 issue）和精度不一致（强制 FP32），按检查清单过一遍，基本能覆盖 90% 的坑。