作者：昇腾实战派

小模型在NPU上的推理部署： 【知识地图】

背景概述

本文旨在介绍如何利用华为昇腾的 Torch Ascend Intermediate Representation (TorchAIR) 技术，在昇腾AI处理器上对PyTorch模型进行图编译优化，以提升推理性能。我们将以Meta的DINOv3模型为例，演示如何使用 torch.compile 结合TorchAIR后端实现"一次编译、多次运行"的高效推理模式。

什么是 TorchAIR？

Torch Ascend Intermediate Representation (TorchAIR) 是华为昇腾为PyTorch开发者提供的高级图编译优化工具。它基于 PyTorch 2.0 引入的 torch.compile API，通过将PyTorch模型的计算图转换为昇腾硬件友好的中间表示（IR），实现更深层次的性能优化。

与基础的 torch_npu 相比，TorchAIR 提供了更激进的优化策略：

• 全图捕获：将整个模型的计算图一次性捕获并优化
• 算子融合：自动识别并融合多个算子，减少内存访问开销
• 内存优化：智能调度内存分配，提升缓存利用率
• 常量折叠：在编译期完成可确定的计算，减少运行时开销

这种"编译-执行"分离的模式特别适合固定模型的批量推理场景，能够显著提升推理吞吐量。

环境准备：版本配套与容器部署

版本配套要求

目前TorchAir暂未提供单独的发布包，而是作为Ascend Extension for PyTorch的三方库，随着torch_npu一起发布。请直接安装torch_npu插件，即可使用TorchAir。

严格核对版本兼容性是成功运行的前提​。您需要确保 ​PyTorch版本​、torch_npu插件版本 与底层的 CANN版本 三者完全匹配。

通常，最省心的方式是直接使用昇腾官方提供的、已配置好环境的Docker镜像。下表列出了部分常见的版本配套关系供参考：

​版本选择逻辑​：通常首先根据您的服务器环境确定 ​CANN 驱动版本​，然后根据上表选择匹配的 PyTorch 和 torch_npu 版本进行安装。完整表格请参阅 torch_npu 官方文档。

拉取并运行Docker容器

使用官方提供的容器环境可以避免复杂的依赖问题：

# 拉取包含TorchAIR的完整环境镜像
docker pull --platform=arm64 swr.cn-south-1.myhuaweicloud.com/ascendhub/mindie:2.1.RC2-300I-Duo-py311-openeuler24.03-lts

# 请替换 ${NAME} 为您的容器名，${IMAGES_ID} 为上一步拉取的镜像ID
docker run --name ${NAME} -it -d --net=host --shm-size=500g \
  --privileged=true \
  -w /home \
  --device=/dev/davinci_manager \
  --device=/dev/hisi_hdc \
  --device=/dev/devmm_svm \
  --entrypoint=bash \
  -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
  -v /usr/local/dcmi:/usr/local/dcmi \
  -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
  -v /usr/local/sbin:/usr/local/sbin \
  -v /home:/home \
  -v /data:/data \
  -v /tmp:/tmp \
  -v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime \
  ${IMAGES_ID}

安装必要依赖

# 升级transformers到与DINOV3匹配版本，以及安装较新的的PyTorch和Torch_npu，可以更好的兼容DINOV3模型
pip install torch==2.8.0 torchaudio==2.8.0 torch_npu==2.8.0 torchvision==0.23.0 transformers==4.56.0

推理实践：使用TorchAIR编译优化DINOv3

以下是完整的代码示例，展示了如何使用TorchAIR对DINOv3模型进行图编译优化和推理：

import torch
import torch_npu
import torchair  # 导入TorchAIR库
from transformers import AutoImageProcessor, AutoModel
from PIL import Image
import numpy as np
# 1. 初始化昇腾设备
def init_npu_device(device_id=0):
    """初始化NPU设备并设置优化选项"""
    # 关闭JIT编译，避免初期调试时的复杂问题
    torch_npu.npu.set_compile_mode(jit_compile=False)
    # 关闭内部格式转换，确保精度与CPU/GPU一致
    torch_npu.npu.config.allow_internal_format = False
    # 指定设备并设置为当前设备
    device = torch.device(f'npu:{device_id}')
    torch.npu.set_device(device)
    print(f"Using device: {device}")
    return device

device = init_npu_device(0)

# 2. 加载模型并迁移至NPU
pretrained_model_name = "/path/to/your/dinov3-vitb16-pretrain-lvd1689m/"  # 请替换为实际路径或Hugging Face模型ID
model = AutoModel.from_pretrained(pretrained_model_name)
model.eval()  # 切换到评估模式
model.to(device)  # 关键步骤：将模型迁移到NPU
print("Model loaded and moved to NPU.")

# 3. 准备输入数据
processor = AutoImageProcessor.from_pretrained(pretrained_model_name)
# 加载一张示例图片
image = Image.open("example.jpg").convert("RGB")
# 处理图片，并确保输入张量也在NPU上
inputs = processor(images=image, return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}  # 关键步骤：将数据迁移到NPU

# 4. 配置TorchAIR编译器
config = torchair.CompilerConfig()

# 5. 获取NPU后端并编译模型
npu_backend = torchair.get_npu_backend(compiler_config=config)
compiled_model = torch.compile(model, backend=npu_backend, fullgraph=True)
print("图编译成功，模型已优化")


# 6. 执行推理（编译后的模型首次运行会较慢，后续运行会更快）
compiled_model.eval()# 4. 执行推理
with torch.no_grad():  # 或 torch.inference_mode()
    outputs = compiled_model(**inputs)

# 7. 处理输出
# 例如，获取CLS Token特征并保存
features = outputs.last_hidden_state[:, 0]  # 取CLS token
features_numpy = features.cpu().numpy()  # 将结果移回CPU并转为NumPy
print(f"Extracted features shape: {features_numpy.shape}")
# 可以继续用于后续任务或保存
# np.save("features.npy", features_numpy)

TorchAIR提供了更多丰富的配置选项，可前往昇腾官方社区了解

总结

本文演示了使用 TorchAIR 在昇腾NPU上对DINOv3模型进行图编译优化的完整流程。整个过程可以概括为：​预处理输入 -> 模型与数据迁移至NPU -> TorchAIR图编译优化 -> 高效推理执行 -> 后处理输出​。

其中最关键的一步，是利用 torch.compile 结合TorchAIR后端，将PyTorch模型的计算图转换为昇腾硬件友好的中间表示并进行深度优化，实现"一次编译、多次运行"的高效推理模式。而通过 .to(device) 这一PyTorch原生语法将模型和输入数据部署到NPU，则体现了昇腾AI处理器生态与PyTorch的无缝兼容性。