小模型在昇腾NPU上的推理部署：【ONNX 模型快速部署】

作者昇腾实战派小模型在NPU上的推理部署【知识地图】背景概述本文为开发者提供一份实用指南演示如何利用ONNX Runtime与华为 CANN的对接能力将通用 ONNX 模型快速部署到昇腾 AI 处理器NPU上进行高效推理。通过标准化的中间格式和高效的执行提供者实现一次导出、多处部署是拥抱国产 AI 硬件生态的便捷路径。什么是 ONNX Runtime CANNONNX Runtime (ORT)是一个跨平台的高性能推理引擎支持多种硬件后端。CANN是华为昇腾 AI 处理器的计算架构。两者的结合点在于ONNX Runtime 的 CANN Execution Provider (EP)。该执行提供者使得 ONNX Runtime 能够将计算图或算子直接调度到昇腾 NPU 上执行。其核心优势在于标准化基于开放的 ONNX 模型格式避免框架锁定高性能利用 CANN 的图编译和算子优化能力易部署无需重新训练只需一次模型转换多后端支持同一套 API 可灵活选择 CPU/GPU/NPU 等不同后端关键前提版本配套与环境要求建议使用 昇腾官方 CANN 镜像 创建容器并安装 onnxruntime 和 onnxruntime-cann。部署实践以 YOLOv5 模型为例环境准备安装 ONNX Runtime CANN建议在昇腾 NPU 开发环境或官方 Docker 容器中进行。# 1. 进入昇腾开发环境假设已安装 CANNsource/usr/local/Ascend/ascend-toolkit/set_env.sh# 2. 安装 ONNX Runtime选择与 CANN 兼容的版本pipinstallonnx1.20.0onnxruntime1.23.2 onnxruntime-cann1.23.2准备 ONNX 模型请前往 Ultralytics 官网 下载 YOLOv5 模型并使用官方脚本导出为 ONNX 格式。使用 ONNX Runtime CANN 进行推理以下为关键步骤的代码示例展示了如何使用 onnxruntime_cann 在昇腾 NPU 上加载模型并进行推理。1. 初始化模型importonnxruntimeasortdefload_model(model_path,device_id):# 设置 NPU 设备torch.npu.set_device(fnpu:{device_id})providers[(CANNExecutionProvider,{device_id:device_id,arena_extend_strategy:kNextPowerOfTwo,npu_mem_limit:4*1024*1024*1024,enable_cann_graph:True,}),]sessionort.InferenceSession(model_path,providersproviders)returnsession2. 预处理importcv2importnumpyasnpimporttorchfromultralytics.data.augmentimportLetterBoxdefpreprocess_image(image_path,target_size(640,640)):original_imagecv2.imread(image_path)# 使用 LetterBox 进行预处理保持长宽比imageLetterBox(target_size,autoFalse)(imageoriginal_image)imageimage.transpose(2,0,1)# HWC to CHWimagenp.ascontiguousarray(image)imagetorch.from_numpy(image).float()/255.0imageimage.unsqueeze(0)# 添加 batch 维度returnimage.numpy(),original_image3. 推理# 加载模型model_pathyolov5s.onnxdevice_id0# NPU 设备 IDsessionload_model(model_path,device_id)# 获取输入输出名称input_namesession.get_inputs()[0].name output_namesession.get_outputs()[0].name# 预处理图像input_image,original_imagepreprocess_image(test.jpg)# 执行推理outputssession.run([output_name],{input_name:input_image})4. 后处理后处理部分通常包括非极大值抑制NMS等操作此处以 YOLOv5 为例进行简单展示fromultralytics.utils.nmsimportnon_max_suppressiondefpostprocess(prediction,original_shape,input_shape):# 使用非极大值抑制prednon_max_suppression(prediction,conf_thres0.25,iou_thres0.45)# 将检测框映射回原图尺寸fordetinpred:ifdetisnotNoneandlen(det):det[:,:4]scale_boxes(input_shape,det[:,:4],original_shape).round()returnpred更多资源更多详细使用教程请参考 ONNX Runtime CANN 官方文档。总结本文演示了使用ONNX Runtime CANN在昇腾 NPU 上进行模型推理的核心流程。通过本文提供的示例代码开发者可以了解如何在昇腾 NPU 上使用 ONNX Runtime CANN 执行提供者加载 ONNX 模型并进行推理。主要步骤包括模型初始化配置 CANNExecutionProvider 并创建推理会话预处理将输入图像转换为模型所需的格式推理调用 ONNX Runtime 的 run 方法在 NPU 上执行计算后处理对模型输出进行解析和转换本文仅提供了关键步骤的代码片段完整的部署方案需要根据具体模型和业务需求进行补充。相比直接使用torch_npuONNX Runtime 方案更适合需要跨框架部署、已有 ONNX 格式模型或追求标准化推理流程的场景。随着 ONNX 生态和 CANN 的不断完善这一方案将成为昇腾 NPU 推理部署的重要选择。