Ascend-SACT/yolo-v8
模型介绍文件和版本Pull Requests讨论分析
下载使用量0

yolov8迁移昇腾NPU操作指导

概述：

YOLOv8 是单阶段实时目标检测框架，以 “Anchor-Free + 解耦头 + 多任务统一” 为核心，兼顾精度、速度与工程易用性，广泛用于工业检测、边缘部署与多场景 CV 任务。

本文介绍在昇腾910B3上基于torch_npu对Yolov8进行训练和推理的全流程。

一、 环境信息

  • 模型：yolov8

  • AI加速卡：910B3

  • CPU架构：ARM

  • CANN：8.5.0

  • numpy：< 2.0

二、 准备流程

1、 容器启动

docker run -it -u root -d --net=host \
--privileged \
--ipc=host \
--device=/dev/davinci_manager \
--device=/dev/devmm_svm \
--device=/dev/hisi_hdc \
-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 /usr/local/Ascend/driver/tools/hccn_tool:/usr/local/Ascend/driver/tools/hccn_tool \
--name yolov8\
quay.io/ascend/cann:8.5.0 \
/bin/bash

说明：后续操作是在容器内

2、 安装依赖

apt update
apt install libxcb1 libx11-xcb1
apt install -y libgl1-mesa-glx
apt install -y libglib2.0-0 libsm6 libxext6 libxrender-dev libgomp1
apt-get install -y zlib1g-dev build-essential
pip install torch==2.1.0
pip install ultralytics==8.2.40
pip install decorator
pip install protobuf
pip install torchvision==0.16.0 torch-npu==2.1.0.post17
pip install attrs
pip install "numpy<2.0"pip3 install -v 'git+https://gitee.com/ascend/tools.git#egg=aclruntime&subdirectory=ais-bench_workload/tool/ais_bench/backend'
pip3 install -v 'git+https://gitee.com/ascend/tools.git#egg=ais_bench&subdirectory=ais-bench_workload/tool/ais_bench'

3、下载数据集

cd /root

mkdir yolov8

cd yolov8

下载训练数据集

https://huggingface.co/datasets/detection-datasets/coco

三、迁移流程

1、训练

参照如下代码编写训练任务。如果是做简单尝试，可以使用coco10这样的小数据集进行验证。

import os
import argparse
import torch
from ultralytics import settings

# 1. 最先执行 NPU 劫持
try:
    import torch_npu
    from torch_npu.contrib import transfer_to_npu
    import torch.distributed as dist
except ImportError:
    pass

from ultralytics import YOLO
from ultralytics.utils import LOGGER, colorstr

# 修改默认下载路径
# /root/.config/Ultralytics/settings.json，缺省配置datasets_dir是root根目录
settings.update({'datasets_dir': '/home/task/datasets'})

def set_npu_env():
    """配置 910B3 最佳性能环境变量"""
    npu_envs = {
        'ASCEND_GLOBAL_LOG_LEVEL': '3',
        'HCCL_CONNECT_TIMEOUT': '1800',
        'PYTORCH_NPU_ALLOC_CONF': 'max_split_size_mb:512',
        'ACL_RTM_ENABLE': '1',
        'TASK_QUEUE_ENABLE': '1'
    }
    os.environ.update(npu_envs)
    if hasattr(torch.npu, 'set_compile_mode'):
        torch.npu.set_compile_mode(jit_compile=False)

def train_yolo(args):
    set_npu_env()

    # 获取分布环境变量
    local_rank = int(os.environ.get('LOCAL_RANK', -1))
    world_size = int(os.environ.get('WORLD_SIZE', 1))

    # 处理设备逻辑
    if local_rank != -1:
        # 多卡模式：手动初始化进程组
        device = f'npu:{local_rank}'
        torch.npu.set_device(device)
        if not dist.is_initialized():
            # 昇腾必须使用 hccl 后端
            dist.init_process_group(backend='hccl', rank=local_rank, world_size=world_size)
        LOGGER.info(f"进程 {local_rank} 启动，共 {world_size} 卡，使用设备: {device}")
    else:
        # 单卡模式
        device = args.device
        LOGGER.info(f"单卡启动，使用设备: {device}")

    # 加载模型
    model = YOLO(args.model)

    # 准备训练参数
    train_params = vars(args).copy()
    
    # 彻底移除可能冲突的 key
    for k in ['val_after_train', 'device']:
        train_params.pop(k, None)

    # 执行训练
    # 关键：在 torchrun 环境下，不显式传 device 参数给 train()，
    # 框架会自动识别已初始化的进程组
    if local_rank != -1:
        model.train(**train_params)
    else:
        model.train(**train_params, device=device)

    # 仅在主进程执行验证
    if args.val_after_train and local_rank in [-1, 0]:
        model.val()

def get_args():
    parser = argparse.ArgumentParser(description='YOLOv12 Ascend 910B3 Optimized')
    parser.add_argument('--model', type=str, default='yolov8n.pt')
    parser.add_argument('--data', type=str, default='coco10/coco10.yaml')
    parser.add_argument('--epochs', type=int, default=100)
    parser.add_argument('--batch', type=int, default=16)
    parser.add_argument('--imgsz', type=int, default=640)
    parser.add_argument('--device', type=str, default='0')
    parser.add_argument('--workers', type=int, default=8)
    parser.add_argument('--amp', action='store_true', default=True)
    parser.add_argument('--val_after_train', action='store_true', default=True)
    parser.add_argument('--project', type=str, default='runs/train-npu')
    parser.add_argument('--name', type=str, default='exp')
    return parser.parse_args()

if __name__ == '__main__':
    args = get_args()
    train_yolo(args)

开始训练：

python train.py --device 0 --batch 16 --data coco10/coco10.yaml --epochs 20 --model yolov8n.pt

训练结果分析：

训练结果的loss曲线位于对应任务文件夹中的result.png（路径可在日志中查看），可结合loss曲线进行训练效果分析

2、 模型转换

（1）pt转换onnx

参考如下代码将pt格式文件转换为onnx，执行python pt2onnx.py

from ultralytics import YOLO

# 修改为训练好的权重文件
model = YOLO("best.pt")

# PT文件输出为onnx文件
model.export(
    format="onnx",
    imgsz=640,           # 指定图片尺寸，可提升atc转换效率
    opset=17,            # ONNX opset 版本，最新默认值
    simplify=True,       # 是否简化模型（调用 onnx-simplifier），True可去除冗余节点，降低atc转换失败概率
    dynamic=False,       # 是否启用动态 shape，默认False
    half=False           # 是否使用 FP16（半精度）
)

（2）onnx转换OM模型

参考如下脚本，在工作目录/home/task生成convert.sh脚本，执行convert.sh。即可在工作目录生成best.om文件。

atc --model=best.onnx \
    --framework=5 \
    --output=best \
    --input_format=NCHW \
    --input_shape="images:1,3,640,640" \
    --soc_version=Ascend910B3 \
    --precision_mode=allow_mix_precision  # 允许推理时混合精度

3、 推理

参考如下脚本，生成推理任务脚本。

# -*- coding: utf-8 -*-
import numpy as np
import cv2
import argparse
from pathlib import Path
from ais_bench.infer.interface import InferSession

# 标准 COCO 80 类别 (确保与你的数据定义一致)
CLASSES = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']

def preprocess(img, sz=(640, 640)):
    """Letterbox 预处理：保持比例缩放并填充背景"""
    h, w = img.shape[:2]
    r = min(sz[0] / h, sz[1] / w)
    new_unpad = int(round(w * r)), int(round(h * r))
    dw, dh = (sz[1] - new_unpad[0]) / 2, (sz[0] - new_unpad[1]) / 2
    img_res = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
    img_pad = cv2.copyMakeBorder(img_res, int(round(dh - 0.1)), int(round(dh + 0.1)), 
                                 int(round(dw - 0.1)), int(round(dw + 0.1)), 
                                 cv2.BORDER_CONSTANT, value=(114, 114, 114))
    # BGR 转 RGB, 归一化, HWC 转 CHW
    tensor = cv2.cvtColor(img_pad, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
    return np.ascontiguousarray(np.transpose(tensor, (2, 0, 1))[None, ...]), r, (dw, dh)

def postprocess(pred, conf_thru, iou_thru):
    """后处理：解析 [84, 8400] 形状的输出"""
    pred = np.squeeze(pred) # [84, 8400]
    boxes = pred[:4, :].T   # [8400, 4] (x,y,w,h)
    scores = np.max(pred[4:, :], axis=0) # [8400]
    class_ids = np.argmax(pred[4:, :], axis=0)
    
    # 阈值过滤
    idx = scores > conf_thru
    boxes, scores, class_ids = boxes[idx], scores[idx], class_ids[idx]
    if len(boxes) == 0: return [], [], []

    # XYWH 转 XYXY
    xyxy = np.copy(boxes)
    xyxy[:, 0], xyxy[:, 1] = boxes[:, 0] - boxes[:, 2]/2, boxes[:, 1] - boxes[:, 3]/2
    xyxy[:, 2], xyxy[:, 3] = boxes[:, 0] + boxes[:, 2]/2, boxes[:, 1] + boxes[:, 3]/2
    
    # NMS (OpenCV 自带的更简洁)
    indices = cv2.dnn.NMSBoxes(xyxy.tolist(), scores.tolist(), conf_thru, iou_thru)
    if len(indices) > 0:
        indices = indices.flatten()
        return xyxy[indices], scores[indices], class_ids[indices]
    return [], [], []

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', '-m', required=True, help='OM model path')
    parser.add_argument('--source', '-s', required=True, help='Image or Directory')
    parser.add_argument('--out', '-o', default='./results', help='Output directory')
    parser.add_argument('--conf', type=float, default=0.25)
    args = parser.parse_args()

    # 初始化
    session = InferSession(device_id=0, model_path=args.model)
    out_path = Path(args.out); out_path.mkdir(exist_ok=True, parents=True)
    src = Path(args.source)
    files = list(src.glob('*.jpg')) if src.is_dir() else [src]

    print(f"Processing {len(files)} images...")
    for f in files:
        im0 = cv2.imread(str(f))
        if im0 is None: continue
        
        # 推理
        tensor, ratio, (dw, dh) = preprocess(im0)
        raw_out = session.infer([tensor])[0]
        boxes, scores, clss = postprocess(raw_out, args.conf, 0.45)

        # 绘图
        for box, score, cls in zip(boxes, scores, clss):
            x1 = int((box[0] - dw) / ratio)
            y1 = int((box[1] - dh) / ratio)
            x2 = int((box[2] - dw) / ratio)
            y2 = int((box[3] - dh) / ratio)
            
            label = f"{CLASSES[int(cls)]} {score:.2f}"
            cv2.rectangle(im0, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(im0, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

        cv2.imwrite(str(out_path / f.name), im0)
    
    print(f"✨ Finished. Results in {args.out}")

if __name__ == '__main__':
    main()

执行推理任务：

python infer_om.py -m best.om -s /root/yolov8/coco10/images -o ./result --conf 0.6

结果：