作者：昇腾实战派

背景概述

本文旨在为开发者提供一份清晰的指南，演示如何利用华为 Ascend Extension for PyTorch (torch_npu)，将经典的PyTorch模型（以Meta的DINOv3为例）快速迁移到昇腾AI处理器上进行高效推理。通过“最小代码改动”实现算力跃迁，是拥抱国产AI硬件生态的便捷路径。

什么是 torch_npu？

Ascend Extension for PyTorch（简称 torch_npu）是华为为昇腾AI处理器开发的官方PyTorch插件。其核心价值在于无缝兼容原生PyTorch生态。

它通过底层的 “monkey-patch”动态替换技术，将PyTorch运算自动对接至昇腾计算库，从而实现了对NPU硬件的原生支持。对于开发者而言，迁移成本极低——在大多数情况下，只需将代码中的设备标识从 "cuda" 改为 "npu"，即可让现有模型在昇腾硬件上执行，大幅降低了使用国产AI算力的门槛。

注意：torch_npu 是一个“桥梁”，其本身不包含计算能力，底层依赖昇腾计算架构 CANN。因此，环境配置中版本匹配至关重要。

关键前提：严格的版本配套

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

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

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

推理实践：以 DINOv3 模型为例

环境准备：拉取并运行Docker容器

建议通过Docker容器来获得一致且干净的环境。

1. 拉取镜像：

   # 请根据您的CANN版本和系统架构（x86或arm）选择合适的镜像
   docker pull --platform=arm64 swr.cn-south-1.myhuaweicloud.com/ascendhub/mindie:2.1.RC2-300I-Duo-py311-openeuler24.03-lts
   

2. 创建并启动容器：

   # 请替换 ${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}
   

   注意：--platform=arm64 参数仅适用于ARM架构服务器。如果您使用的是x86服务器，请从昇腾社区镜像仓库查找并拉取对应的x86镜像。

容器内安装依赖

进入容器后，安装匹配版本的PyTorch套件及模型库。

# 升级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

使用 torch_npu 加载模型并推理

以下是一个完整的示例脚本，展示了从加载到推理的全过程。

import torch
import torch_npu
from transformers import AutoImageProcessor, AutoModel
from PIL import Image

# 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. 执行推理
with torch.no_grad():  # 或 torch.inference_mode()
    outputs = model(**inputs)

# 5. 处理输出
# 例如，获取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)

总结

本文演示了使用 torch_npu 在昇腾NPU上进行小模型推理的核心流程。整个过程可以概括为：预处理输入 -> 模型与数据迁移至NPU -> 执行推理 -> 后处理输出。

其中最关键的一步，便是利用 .to(device) 这一PyTorch原生语法，将模型和输入数据精准地部署到NPU设备上。这种极简的迁移方式，正是 torch_npu 插件强大兼容性的体现。

希望这篇指南能帮助您顺利踏上昇腾NPU的推理加速之旅。