我是提前训练好了YOLO模型,pt文件的。现在要在华为加速卡环境下部署推理。

0. 准备工作:先根据官方文档,把开发环境装好了

CANN安装部署-昇腾社区

 

1、ONNX模型转OM模型(通过昇腾张量编译器ATC编译出来的)

由于我之前用的是pt模型,现在需要先将pt转为ONNX,再把ONNX转为OM模型。

转换代码如下,直接用YOLO的转换方法

from ultralytics import YOLO
from pathlib import Path

WEIGHTS_DIR = Path(__file__).resolve().parent / "weights"
ONNX_DIR = Path(__file__).resolve().parent / "onnx"

def pt2onnx(pt_path: str, onnx_path: str):
    model = YOLO(pt_path)
    model.export(format="onnx", imgsz=640, opset=11, simplify=True)
    print(f"ONNX model saved to: {onnx_path}")

if __name__ == "__main__":
    ONNX_DIR.mkdir(exist_ok=True)
    pt_path = WEIGHTS_DIR / "best.pt"
    onnx_path = ONNX_DIR / "best.onnx"
    if not pt_path.exists():
        print(f"Error: {pt_path} not found")
    else:
        pt2onnx(str(pt_path), str(onnx_path))

ONXX转OM模型zh

去服务器上执行转换命令

# 3. 执行 ATC 转换命令 
# --soc_version 需替换为实际芯片型号,
# --input_shape 需根据模型实际输入定义
# --framework=5:代表 ONNX 模型。

atc --framework=5 \
    --model=./best.onnx \
    --output=./model_bs1 \
    --input_format=NCHW \
    --input_shape="images:1,3,640,640" \
    --soc_version=Ascend310P3 \
    --log=error

过程运行的时候,会发现可能缺少一些依赖,可以执行下面的命令安装

第一次运行:pip3 install numpy
第二次运行:pip3 install sympy attrs psutil scipy

转换完成后,当前目录就出现了  om 模型文件

2. 推理运行

以下为官方推理demo,放到服务器上。执行。

函数:create_tensor_desc-CANN商用版9.0.0-昇腾社区 这里是api文档

"""
Ascend 昇腾推理脚本 - YOLO11 目标检测
基于 pyACL (Python Ascend Computing Language) 开发,用于在昇腾 NPU 上执行 YOLO 推理。

使用方法:
    python3 ascend_infer.py --model ./model_bs1.om --image ./fire11.jpg --output result.jpg

前置条件:
    1. 已通过 ATC 工具将 ONNX 模型转换为 .om 格式
    2. 已安装 pyACL、OpenCV、NumPy
    3. 已设置 Ascend 环境变量 (source set_env.sh)
"""

import os
import sys
import numpy as np
import cv2

# 导入 pyACL 模块(仅在昇腾设备上可用)
try:
    import acl
except ImportError:
    print("pyACL not found. Please run on Ascend device.")
    sys.exit(1)

# ==================== 常量定义 ====================
# 内存分配策略:优先分配大页内存,提高大块内存分配效率
ACL_MEM_MALLOC_HUGE_FIRST = 0

# 内存拷贝方向
ACL_MEMCPY_HOST_TO_DEVICE = 1  # 主机(Host) -> 设备(Device),用于上传输入数据
ACL_MEMCPY_DEVICE_TO_HOST = 2  # 设备(Device) -> 主机(Host),用于下载输出结果


class YOLOInference:
    """
    YOLO 推理类,封装了昇腾 NPU 上完整的推理流程:
    1. 初始化 AscendCL 环境
    2. 加载 .om 模型
    3. 准备输入输出数据集(Device 内存)
    4. 图像预处理(Resize、归一化、格式转换)
    5. 执行推理
    6. 后处理(坐标还原、NMS、阈值过滤)
    7. 资源释放
    """

    def __init__(self, model_path, input_size=640, conf_thresh=0.25, iou_thresh=0.45):
        """
        初始化推理参数(尚未初始化 AscendCL,需调用 init())

        参数:
            model_path: .om 模型文件路径
            input_size: 模型输入图像尺寸(正方形),默认 640x640
            conf_thresh: 置信度阈值,低于此值的检测框将被过滤,默认 0.25
            iou_thresh: NMS 的 IoU 阈值,重叠度高于此值的框将被抑制,默认 0.45
        """
        self.model_path = model_path
        self.input_size = input_size
        self.conf_thresh = conf_thresh
        self.iou_thresh = iou_thresh
        self.device_id = 0  # 使用第 0 号昇腾设备

        # 以下变量在后续步骤中初始化
        self.model_id = None        # 模型 ID(由 load_from_file 返回)
        self.model_desc = None      # 模型描述符(包含输入输出信息)
        self.input_dataset = None   # 输入数据集(包含输入 buffer)
        self.output_dataset = None  # 输出数据集(包含输出 buffer)
        self.input_data = None      # 输入数据列表,每个元素包含 buffer 地址、大小等
        self.output_data = None     # 输出数据列表,每个元素包含 buffer 地址、大小等
        self.output_dims = None     # 输出 tensor 形状,如 [1, 6, 8400]

    # ==================== 第一步:初始化环境与资源 ====================
    def init(self):
        """
        初始化 AscendCL 运行环境:
        1. acl.init() - 初始化 AscendCL 框架
        2. acl.rt.set_device() - 绑定指定的昇腾 NPU 设备
        """
        # 初始化 AscendCL,必须第一个调用
        ret = acl.init()
        if ret != 0:
            raise RuntimeError(f"acl.init failed: {ret}")

        # 绑定设备,后续所有操作都在该设备上执行
        ret = acl.rt.set_device(self.device_id)
        if ret != 0:
            raise RuntimeError(f"acl.rt.set_device failed: {ret}")

    # ==================== 第二步:加载 OM 模型 ====================
    def load_model(self):
        """
        加载 .om 模型并获取模型信息:
        1. 从文件加载模型,获得 model_id
        2. 创建模型描述符,获取输入输出的维度、大小等信息
        3. 在 Device 上申请输入输出内存
        4. 创建输入输出数据集(Dataset + DataBuffer)
        """
        # 从 .om 文件加载模型到设备
        self.model_id, ret = acl.mdl.load_from_file(self.model_path)
        if ret != 0:
            raise RuntimeError(f"load_from_file failed: {ret}")

        # 创建模型描述符并填充模型信息
        self.model_desc = acl.mdl.create_desc()
        ret = acl.mdl.get_desc(self.model_desc, self.model_id)
        if ret != 0:
            raise RuntimeError(f"get_desc failed: {ret}")

        # 获取输出 tensor 形状,如 [1, 6, 8400]
        # get_output_dims 返回格式: ({'name': '...', 'dims': [1,6,8400]}, 0)
        out_dims_info = acl.mdl.get_output_dims(self.model_desc, 0)
        self.output_dims = list(out_dims_info[0]["dims"])
        print(f"Output shape: {self.output_dims}")

        # 创建输入和输出数据集,包含 Device 内存分配和 DataBuffer 创建
        self.input_dataset, self.input_data = self._prepare_dataset("input")
        self.output_dataset, self.output_data = self._prepare_dataset("output")

    def _prepare_dataset(self, io_type):
        """
        准备输入或输出数据集(核心函数)

        流程:
        1. 获取输入/输出的数量和各自的 buffer 大小
        2. 为每个输入/输出在 Device 上申请内存 (acl.rt.malloc)
        3. 将 Device 内存地址包装为 DataBuffer (acl.create_data_buffer)
        4. 将 DataBuffer 加入 Dataset (acl.mdl.add_dataset_buffer)

        参数:
            io_type: "input" 或 "output"

        返回:
            (dataset, datas):
                dataset: acl Dataset 对象,传给 mdl.execute 使用
                datas: 列表,每个元素包含 buffer 地址、DataBuffer 句柄、大小
        """
        # 根据输入/输出类型,获取对应的数量和大小查询函数
        if io_type == "input":
            io_num = acl.mdl.get_num_inputs(self.model_desc)      # 输入数量(YOLO 通常为 1)
            get_size = acl.mdl.get_input_size_by_index            # 获取指定输入的字节数
        else:
            io_num = acl.mdl.get_num_outputs(self.model_desc)     # 输出数量(YOLO 通常为 1)
            get_size = acl.mdl.get_output_size_by_index           # 获取指定输出的字节数

        # 创建数据集容器
        dataset = acl.mdl.create_dataset()
        datas = []

        for i in range(io_num):
            # 获取当前输入/输出的 buffer 大小(字节)
            # 例如输入: 1*3*640*640*4 = 4915200 字节 (float32)
            # 例如输出: 1*6*8400*4 = 201600 字节 (float32)
            buffer_size = get_size(self.model_desc, i)

            # 在 Device(NPU)上申请内存
            # 返回: (device_ptr, error_code)
            # ACL_MEM_MALLOC_HUGE_FIRST 表示优先分配大页内存
            buffer, ret = acl.rt.malloc(buffer_size, ACL_MEM_MALLOC_HUGE_FIRST)
            if ret != 0:
                raise RuntimeError(f"rt.malloc failed: {ret}")

            # 将 Device 内存地址包装为 DataBuffer 对象
            # DataBuffer 是 pyACL 用于描述一块内存地址和大小的句柄
            data_buffer = acl.create_data_buffer(buffer, buffer_size)

            # 将 DataBuffer 加入到 Dataset 中
            # Dataset 是 mdl.execute 需要的输入/输出容器
            _, ret = acl.mdl.add_dataset_buffer(dataset, data_buffer)
            if ret != 0:
                raise RuntimeError(f"add_dataset_buffer failed: {ret}")

            # 保存信息,后续 memcpy 和释放时需要用到
            datas.append({
                "buffer": buffer,       # Device 内存地址(用于 memcpy)
                "data": data_buffer,    # DataBuffer 句柄(用于销毁)
                "size": buffer_size     # buffer 大小(字节)
            })

            if io_type == "input":
                print(f"Input {i}: {buffer_size} bytes")
            else:
                print(f"Output {i}: {buffer_size} bytes")

        return dataset, datas

    # ==================== 第三步:图像预处理 ====================
    def preprocess(self, image_path):
        """
        将原始图像预处理为模型输入格式:

        处理步骤:
        1. 读取图像(OpenCV 默认 BGR 格式)
        2. BGR -> RGB 通道转换
        3. Resize 到模型输入尺寸(如 640x640)
        4. 像素值归一化到 [0, 1](除以 255)
        5. HWC -> CHW 维度转换(Height,Width,Channel -> Channel,Height,Width)
        6. 增加 batch 维度(NCHW 格式)

        参数:
            image_path: 输入图像文件路径

        返回:
            (input_data, orig_img, orig_w, orig_h):
                input_data: 预处理后的 numpy 数组,shape=(1,3,640,640),dtype=float32
                orig_img: 原始图像(用于绘制结果)
                orig_w: 原始图像宽度
                orig_h: 原始图像高度
        """
        # 读取图像,OpenCV 默认 BGR 格式
        img = cv2.imread(image_path)
        if img is None:
            raise FileNotFoundError(f"Cannot read image: {image_path}")

        # 记录原始尺寸,后续后处理时需要将坐标映射回原图
        orig_h, orig_w = img.shape[:2]

        # BGR -> RGB(YOLO 训练时使用 RGB 格式)
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # Resize 到模型输入尺寸(640x640)
        img_resized = cv2.resize(img_rgb, (self.input_size, self.input_size))

        # 归一化到 [0, 1](训练时同样做了此操作)
        img_float = img_resized.astype(np.float32) / 255.0

        # HWC -> CHW:将通道维移到最前面
        # 原始: (640, 640, 3) -> 转置后: (3, 640, 640)
        img_nchw = img_float.transpose(2, 0, 1)[np.newaxis, ...]

        # 增加 batch 维度:(3, 640, 640) -> (1, 3, 640, 640)
        # 返回 float32,因为模型输入是 float32
        return img_nchw.astype(np.float32), img, orig_w, orig_h

    # ==================== 第四步:执行推理 ====================
    def forward(self, input_data):
        """
        在昇腾 NPU 上执行模型推理(同步):

        流程:
        1. 将预处理后的数据从 Host 拷贝到 Device(输入)
        2. 调用 mdl.execute 执行推理
        3. 将推理结果从 Device 拷贝回 Host(输出)

        参数:
            input_data: 预处理后的输入数据,shape=(1,3,640,640),dtype=float32

        返回:
            output: 模型输出,shape=(1,6,8400),dtype=float32
                    6 = 4(坐标: cx,cy,w,h) + 2(类别置信度: fire,smoke)
                    8400 = 检测框数量(不同尺度的特征图拼接后)
        """
        # ========== 1. Host -> Device:将输入数据上传到 NPU ==========
        # 将 numpy 数组转为字节流
        bytes_data = input_data.tobytes()
        # 将字节流转为 pyACL 可识别的指针地址
        bytes_ptr = acl.util.bytes_to_ptr(bytes_data)

        # 执行内存拷贝:从 Host(主机内存)拷贝到 Device(NPU 显存)
        # 参数:目标地址, 目标大小, 源地址, 源大小, 拷贝方向
        ret = acl.rt.memcpy(
            self.input_data[0]["buffer"],      # Device 上的输入 buffer 地址
            self.input_data[0]["size"],        # Device buffer 大小
            bytes_ptr,                         # Host 上的输入数据指针
            len(bytes_data),                   # 数据字节数
            ACL_MEMCPY_HOST_TO_DEVICE          # 拷贝方向:Host -> Device
        )
        if ret != 0:
            raise RuntimeError(f"memcpy H2D failed: {ret}")

        # ========== 2. 执行推理 ==========
        # 调用 NPU 执行模型推理(同步阻塞,直到推理完成)
        # 参数:模型ID, 输入数据集, 输出数据集
        ret = acl.mdl.execute(self.model_id, self.input_dataset, self.output_dataset)
        if ret != 0:
            raise RuntimeError(f"mdl.execute failed: {ret}")

        # ========== 3. Device -> Host:将输出数据下载到主机 ==========
        # 在 Host 上申请临时内存,用于接收输出数据
        buffer_host, ret = acl.rt.malloc_host(self.output_data[0]["size"])
        if ret != 0:
            raise RuntimeError(f"malloc_host failed: {ret}")

        # 执行内存拷贝:从 Device(NPU 显存)拷贝到 Host(主机内存)
        ret = acl.rt.memcpy(
            buffer_host,                       # Host 上的目标地址
            self.output_data[0]["size"],       # 目标大小
            self.output_data[0]["buffer"],     # Device 上的输出 buffer 地址
            self.output_data[0]["size"],       # 源大小
            ACL_MEMCPY_DEVICE_TO_HOST          # 拷贝方向:Device -> Host
        )
        if ret != 0:
            raise RuntimeError(f"memcpy D2H failed: {ret}")

        # 将 Host 内存指针转为字节流,再转为 numpy 数组
        bytes_out = acl.util.ptr_to_bytes(buffer_host, self.output_data[0]["size"])
        # 按模型输出形状 reshape,dtype=float32
        output = np.frombuffer(bytes_out, dtype=np.float32).reshape(self.output_dims)

        # 释放 Host 上的临时内存
        acl.rt.free_host(buffer_host)

        return output

    # ==================== 第五步:后处理 ====================
    def postprocess(self, output, orig_w, orig_h):
        """
        对模型原始输出进行后处理,得到最终检测结果:

        处理步骤:
        1. 解析输出 tensor,分离边界框坐标和类别分数
        2. 置信度过滤:去掉低置信度的检测框
        3. 坐标还原:从模型输入尺寸映射回原始图像尺寸
        4. NMS(非极大值抑制):去除重叠度高的冗余框

        YOLO 输出格式:(1, 6, 8400)
            - 第 0~3 维:cx, cy, w, h(中心点坐标和宽高,归一化到 0~640)
            - 第 4~5 维:fire 置信度, smoke 置信度
            - 8400:检测框总数(不同尺度特征图的锚框数量之和)

        参数:
            output: 模型输出,shape=(1,6,8400)
            orig_w: 原始图像宽度
            orig_h: 原始图像高度

        返回:
            (boxes, scores, class_ids):
                boxes: 检测框坐标 [[x1,y1,x2,y2], ...],已映射到原图尺寸
                scores: 置信度分数
                class_ids: 类别 ID(0=fire, 1=smoke)
        """
        # 去掉 batch 维度:(1,6,8400) -> (6,8400)
        output = output[0]

        # 如果输出格式是 (6, 8400),转置为 (8400, 6) 方便按行处理
        if output.shape[0] == 6:
            output = output.T

        # 分离坐标和分数
        boxes = output[:, :4]     # 前 4 列:cx, cy, w, h
        scores = output[:, 4:]    # 后 2 列:fire_score, smoke_score

        # 取每个检测框的最大类别分数及其类别 ID
        class_ids = np.argmax(scores, axis=1)  # 每行最大值的索引(0 或 1)
        max_scores = np.max(scores, axis=1)    # 每行的最大值

        # 置信度过滤:去掉分数低于阈值的框
        mask = max_scores > self.conf_thresh
        boxes = boxes[mask]
        max_scores = max_scores[mask]
        class_ids = class_ids[mask]

        # 没有检测到目标,返回空数组
        if len(boxes) == 0:
            return np.array([]), np.array([]), np.array([])

        # 坐标转换:从 (cx, cy, w, h) 转为 (x1, y1, x2, y2),并映射回原图尺寸
        # 公式:x1 = (cx - w/2) / input_size * orig_w
        x1 = (boxes[:, 0] - boxes[:, 2] / 2) / self.input_size * orig_w
        y1 = (boxes[:, 1] - boxes[:, 3] / 2) / self.input_size * orig_h
        x2 = (boxes[:, 0] + boxes[:, 2] / 2) / self.input_size * orig_w
        y2 = (boxes[:, 1] + boxes[:, 3] / 2) / self.input_size * orig_h

        # 拼接为 (N, 4) 的 xyxy 格式
        boxes_xyxy = np.stack([x1, y1, x2, y2], axis=1)

        # NMS(非极大值抑制):去除高度重叠的冗余检测框
        # OpenCV 的 NMSBoxes 要求输入为 list 格式
        indices = cv2.dnn.NMSBoxes(
            boxes_xyxy.tolist(),    # 检测框列表
            max_scores.tolist(),    # 对应分数列表
            self.conf_thresh,       # 置信度阈值
            self.iou_thresh         # IoU 阈值(重叠度高于此值则抑制)
        )

        # 处理 NMS 返回结果
        if len(indices) == 0:
            return np.array([]), np.array([]), np.array([])

        if isinstance(indices, np.ndarray):
            indices = indices.flatten()
        elif isinstance(indices, list) and len(indices) > 0:
            indices = np.array(indices).flatten()
        else:
            return np.array([]), np.array([]), np.array([])

        # 返回 NMS 后保留的检测结果
        return boxes_xyxy[indices], max_scores[indices], class_ids[indices]

    # ==================== 绘制检测结果 ====================
    def draw_results(self, image, boxes, scores, class_ids, class_names=None):
        """
        在图像上绘制检测框和标签

        参数:
            image: 原始图像
            boxes: 检测框坐标 [[x1,y1,x2,y2], ...]
            scores: 置信度分数
            class_ids: 类别 ID
            class_names: 类别名称字典,默认 {0: "fire", 1: "smoke"}

        返回:
            img: 绘制了检测框的图像
        """
        img = image.copy()

        # 默认类别名称(火灾检测场景)
        if class_names is None:
            class_names = {0: "fire", 1: "smoke"}

        # 不同类别的颜色(BGR 格式)
        colors = [(0, 0, 255), (0, 165, 255)]  # 红色=fire, 橙色=smoke

        # 没有检测结果,直接返回原图
        if len(boxes) == 0:
            return img

        # 遍历每个检测框,绘制矩形和标签
        for box, score, cls_id in zip(boxes, scores, class_ids):
            x1, y1, x2, y2 = map(int, box)
            color = colors[int(cls_id) % len(colors)]
            name = class_names.get(int(cls_id), f"class_{cls_id}")

            # 绘制检测框(矩形)
            cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)

            # 绘制标签背景和文字
            label = f"{name} {score:.2f}"
            (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1)
            cv2.rectangle(img, (x1, y1 - th - 6), (x1 + tw, y1), color, -1)  # 背景矩形
            cv2.putText(img, label, (x1, y1 - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)

        return img

    # ==================== 第六步:释放资源 ====================
    def destroy(self):
        """
        按相反顺序释放所有 AscendCL 资源:

        释放顺序(与初始化顺序相反):
        1. 销毁 DataBuffer,释放 Device 内存
        2. 销毁 Dataset
        3. 销毁模型描述符
        4. 卸载模型
        5. 重置设备
        6. AscendCL 去初始化
        """
        # 释放输入和输出的 DataBuffer 和 Device 内存
        for dataset_data in [self.input_data, self.output_data]:
            if dataset_data:
                for item in dataset_data:
                    acl.destroy_data_buffer(item["data"])  # 销毁 DataBuffer 句柄
                    acl.rt.free(item["buffer"])            # 释放 Device 内存

        # 销毁 Dataset
        if self.input_dataset:
            acl.mdl.destroy_dataset(self.input_dataset)
        if self.output_dataset:
            acl.mdl.destroy_dataset(self.output_dataset)

        # 销毁模型描述符
        if self.model_desc:
            acl.mdl.destroy_desc(self.model_desc)

        # 卸载模型
        if self.model_id is not None:
            acl.mdl.unload(self.model_id)

        # 重置设备并去初始化
        acl.rt.reset_device(self.device_id)
        acl.finalize()


def main():
    """
    主函数:解析命令行参数,执行完整的推理流程
    """
    import argparse

    parser = argparse.ArgumentParser(description="YOLO11 Fire Detection on Ascend")
    parser.add_argument("--model", type=str, required=True, help="Path to .om model file")
    parser.add_argument("--image", type=str, required=True, help="Path to input image")
    parser.add_argument("--output", type=str, default="result.jpg", help="Output image path")
    parser.add_argument("--conf", type=float, default=0.25, help="Confidence threshold")
    parser.add_argument("--iou", type=float, default=0.45, help="IoU threshold for NMS")
    parser.add_argument("--input-size", type=int, default=640, help="Model input size")
    args = parser.parse_args()

    # 创建推理实例
    detector = YOLOInference(
        model_path=args.model,
        input_size=args.input_size,
        conf_thresh=args.conf,
        iou_thresh=args.iou
    )

    try:
        # 第一步:初始化 AscendCL 环境
        print("Initializing Ascend environment...")
        detector.init()

        # 第二步:加载 .om 模型
        print(f"Loading model: {args.model}")
        detector.load_model()

        # 第三步:图像预处理
        print(f"Preprocessing image: {args.image}")
        input_data, orig_img, orig_w, orig_h = detector.preprocess(args.image)

        # 第四步:执行推理
        print("Running inference...")
        output = detector.forward(input_data)

        # 第五步:后处理
        print("Postprocessing results...")
        boxes, scores, class_ids = detector.postprocess(output, orig_w, orig_h)

        # 打印检测结果
        print(f"Detected {len(boxes)} objects")
        for i, (box, score, cls_id) in enumerate(zip(boxes, scores, class_ids)):
            print(f"  [{i}] class={int(cls_id)}, score={score:.4f}, box={box.astype(int).tolist()}")

        # 绘制结果并保存
        result_img = detector.draw_results(orig_img, boxes, scores, class_ids)
        cv2.imwrite(args.output, result_img)
        print(f"Result saved to: {args.output}")

    finally:
        # 第六步:释放所有资源(无论是否出错都会执行)
        detector.destroy()
        print("Resources released.")


if __name__ == "__main__":
    main()

执行推理

python3 ascend_infer.py --model ./model_bs1.om --image ./fire11.jpg

图片下载下来查看,也是推理成功了,这次的推理用到的就是华为的NPU了

Logo

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。

更多推荐