小模型在昇腾NPU上的推理部署：【使用TorchAIR进行模型图编译推理优化】

作者昇腾实战派小模型在NPU上的推理部署【知识地图】背景概述本文旨在介绍如何利用华为昇腾的Torch Ascend Intermediate Representation (TorchAIR)技术在昇腾AI处理器上对PyTorch模型进行图编译优化以提升推理性能。我们将以Meta的DINOv3模型为例演示如何使用torch.compile结合TorchAIR后端实现一次编译、多次运行的高效推理模式。什么是 TorchAIRTorch Ascend Intermediate Representation (TorchAIR)是华为昇腾为PyTorch开发者提供的高级图编译优化工具。它基于PyTorch 2.0引入的torch.compileAPI通过将PyTorch模型的计算图转换为昇腾硬件友好的中间表示IR实现更深层次的性能优化。与基础的torch_npu相比TorchAIR 提供了更激进的优化策略全图捕获将整个模型的计算图一次性捕获并优化算子融合自动识别并融合多个算子减少内存访问开销内存优化智能调度内存分配提升缓存利用率常量折叠在编译期完成可确定的计算减少运行时开销这种编译-执行分离的模式特别适合固定模型的批量推理场景能够显著提升推理吞吐量。环境准备版本配套与容器部署版本配套要求目前TorchAir暂未提供单独的发布包而是作为Ascend Extension for PyTorch的三方库随着torch_npu一起发布。请直接安装torch_npu插件即可使用TorchAir。严格核对版本兼容性是成功运行的前提​。您需要确保 ​PyTorch版本​、torch_npu插件版本与底层的CANN版本三者完全匹配。通常最省心的方式是直接使用昇腾官方提供的、已配置好环境的Docker镜像。下表列出了部分常见的版本配套关系供参考CANN 版本支持的 PyTorch 版本对应的 torch_npu 版本CANN 8.3.RC12.8.02.8.02.7.12.7.12.6.02.6.0.post3CANN 8.2.RC12.6.02.6.02.5.12.5.1.post1CANN 8.1.RC12.5.12.5.12.4.02.4.0.post4​版本选择逻辑​通常首先根据您的服务器环境确定 ​CANN 驱动版本​然后根据上表选择匹配的 PyTorch 和 torch_npu 版本进行安装。完整表格请参阅 torch_npu 官方文档。拉取并运行Docker容器使用官方提供的容器环境可以避免复杂的依赖问题# 拉取包含TorchAIR的完整环境镜像dockerpull--platformarm64 swr.cn-south-1.myhuaweicloud.com/ascendhub/mindie:2.1.RC2-300I-Duo-py311-openeuler24.03-lts# 请替换 ${NAME} 为您的容器名${IMAGES_ID} 为上一步拉取的镜像IDdockerrun--name${NAME}-it-d--nethost --shm-size500g\--privilegedtrue\-w/home\--device/dev/davinci_manager\--device/dev/hisi_hdc\--device/dev/devmm_svm\--entrypointbash\-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模型pipinstalltorch2.8.0torchaudio2.8.0torch_npu2.8.0torchvision0.23.0transformers4.56.0推理实践使用TorchAIR编译优化DINOv3以下是完整的代码示例展示了如何使用TorchAIR对DINOv3模型进行图编译优化和推理importtorchimporttorch_npuimporttorchair# 导入TorchAIR库fromtransformersimportAutoImageProcessor,AutoModelfromPILimportImageimportnumpyasnp# 1. 初始化昇腾设备definit_npu_device(device_id0):初始化NPU设备并设置优化选项# 关闭JIT编译避免初期调试时的复杂问题torch_npu.npu.set_compile_mode(jit_compileFalse)# 关闭内部格式转换确保精度与CPU/GPU一致torch_npu.npu.config.allow_internal_formatFalse# 指定设备并设置为当前设备devicetorch.device(fnpu:{device_id})torch.npu.set_device(device)print(fUsing device:{device})returndevice deviceinit_npu_device(0)# 2. 加载模型并迁移至NPUpretrained_model_name/path/to/your/dinov3-vitb16-pretrain-lvd1689m/# 请替换为实际路径或Hugging Face模型IDmodelAutoModel.from_pretrained(pretrained_model_name)model.eval()# 切换到评估模式model.to(device)# 关键步骤将模型迁移到NPUprint(Model loaded and moved to NPU.)# 3. 准备输入数据processorAutoImageProcessor.from_pretrained(pretrained_model_name)# 加载一张示例图片imageImage.open(example.jpg).convert(RGB)# 处理图片并确保输入张量也在NPU上inputsprocessor(imagesimage,return_tensorspt)inputs{k:v.to(device)fork,vininputs.items()}# 关键步骤将数据迁移到NPU# 4. 配置TorchAIR编译器configtorchair.CompilerConfig()# 5. 获取NPU后端并编译模型npu_backendtorchair.get_npu_backend(compiler_configconfig)compiled_modeltorch.compile(model,backendnpu_backend,fullgraphTrue)print(图编译成功模型已优化)# 6. 执行推理编译后的模型首次运行会较慢后续运行会更快compiled_model.eval()# 4. 执行推理withtorch.no_grad():# 或 torch.inference_mode()outputscompiled_model(**inputs)# 7. 处理输出# 例如获取CLS Token特征并保存featuresoutputs.last_hidden_state[:,0]# 取CLS tokenfeatures_numpyfeatures.cpu().numpy()# 将结果移回CPU并转为NumPyprint(fExtracted 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的无缝兼容性。