【DCU推理加速】YOLO26 + MIGraphX 8卡并行推理性能测试

前言

YOLO26 是 Ultralytics 于2025年9月发布的最新一代实时目标检测模型,相比 YOLO11 带来了多项架构革新:原生 NMS-Free 端到端推理(默认一对一检测头,无需 NMS 后处理)、移除 DFL(降低检测头复杂度、简化导出)、以及全新的 MuSGD 优化器 + Progressive Loss 训练配方。这些特性使 YOLO26 在部署端天然友好——更简洁的计算图结构有利于推理引擎优化。

本文在海光 BW1100 平台上,使用 MIGraphX 推理引擎对 YOLO26s 模型进行 FP16 编译,并通过 8 卡并行推理脚本测试 COCO 数据集上的吞吐性能。

整体流程:拉取镜像 → 安装依赖 → 原生推理验证 → 导出ONNX → 编译MXR → 8卡并行推理

目录


MIGraphX 简介

MIGraphX 是 AMD ROCm 生态下的高性能深度学习推理引擎,在 DCU(海光加速计算单元)平台上扮演着类似 NVIDIA TensorRT 的角色——将训练好的深度学习模型编译优化为可在 GPU/DCU 上高效执行的推理程序。

核心能力

能力说明
ONNX 解析直接读取标准 ONNX 模型,无需手动实现算子映射
图优化常量折叠、算子融合、死代码消除、内存复用等编译期优化
FP16 量化编译阶段自动将 FP32 算子转换为 FP16,减少显存占用和计算延迟
离线编译 (.mxr)将优化后的计算图序列化为 .mxr 离线文件,运行时直接加载,跳过编译开销
多卡支持通过 device_id 参数指定 GPU,支持多卡并行推理
NHWC 布局优化MIGRAPHX_ENABLE_NHWC=1 自动将 NCHW 转为 NHWC 布局,提升卷积算子访存效率

工作流程

ONNX 模型 → migraphx-driver compile → .mxr 离线文件 → migraphx.load() → 推理执行
              ↑ 图优化 + FP16 量化           ↑ 零编译开销加载

与同类工具的对比:

特性MIGraphXTensorRTONNX Runtime
目标平台AMD GPU / 海光 DCUNVIDIA GPU跨平台
输入格式ONNXONNX / TRTONNX
离线格式.mxr.engine无(JIT)
FP16 量化编译期编译期运行期
开源协议MIT闭源MIT

在 DCU 部署中的定位:DTK(DCU Toolkit)提供了 HIP 运行时和 ROCm 库栈,MIGraphX 则在其之上构建了端到端的推理编译能力。对于需要将 PyTorch 训练模型快速部署到 DCU 的场景,MIGraphX 是首选方案——只需 导出 ONNX → 编译 MXR → 加载推理 三步即可完成部署。


测试环境

项目配置
硬件平台海光BW1100(8 × BW1101 gfx938)
操作系统Ubuntu 22.04
容器镜像pytorch:2.9.0-ubuntu22.04-dtk26.04-py3.10
DTK版本26.04
推理引擎MIGraphX 5.2.0
模型YOLO26s(9.5M params, 20.7 GFLOPs, 48.6 mAP)
数据集COCO Val2017(5000张)

一、环境准备

1. 拉取镜像并创建容器

docker pull harbor.sourcefind.cn:5443/dcu/admin/base/pytorch:2.9.0-ubuntu22.04-dtk26.04-py3.10
docker run -id \
    --shm-size 256g \
    --network=host \
    --name=yolo \
    --privileged \
    --device=/dev/kfd \
    --device=/dev/dri \
    --device=/dev/mkfd \
    --ipc=host \
    --group-add video \
    --cap-add=SYS_PTRACE \
    --security-opt seccomp=unconfined \
    -u root --ulimit stack=-1:-1 --ulimit memlock=-1:-1 \
    -v /opt/hyhal/:/opt/hyhal/:ro \
    -v /data:/data \
    harbor.sourcefind.cn:5443/dcu/admin/base/pytorch:2.9.0-ubuntu22.04-dtk26.04-py3.10 /bin/bash

关键参数说明:

参数作用
--shm-size 256g扩大共享内存,满足多进程大数据交换需求
--device=/dev/kfd --device=/dev/dri --device=/dev/mkfd映射DCU设备节点
--ipc=host共享主机IPC命名空间,支持多进程共享内存
--ulimit stack=-1:-1 --ulimit memlock=-1:-1取消栈大小和内存锁定限制
-v /opt/hyhal/:/opt/hyhal/:ro只读挂载海光HAL库

进入容器:

docker exec -it yolo bash

配置DTK环境:

source /opt/dtk-26.04/env.sh

2. 安装Python依赖

pip config set global.index-url https://pypi.mirrors.ustc.edu.cn/simple/
pip config set global.trusted-host pypi.mirrors.ustc.edu.cn
pip install ultralytics

使用中科大镜像源加速下载。ultralytics 提供 YOLO26 模型加载与导出功能,需确保版本 ≥ 8.3.0 以支持 YOLO26。

3. 安装MIGraphX

wget https://download.sourcefind.cn:65024/file/4/migraphx/DAS1.8/migraphx-5.2.0+das.opt1.dtk2604-cp310-cp310-manylinux_2_28_x86_64.run
bash migraphx-5.2.0+das.opt1.dtk2604-cp310-cp310-manylinux_2_28_x86_64.run

安装完成后验证:

migraphx-driver --version

4. 下载COCO数据集

# 下载 val2017 图片(约1GB,5000张)
wget http://images.cocodataset.org/zips/val2017.zip
unzip val2017.zip

# 下载标注文件并转换为YOLO格式
wget http://images.cocodataset.org/annotations/annotations_trainval2017.zip
unzip annotations_trainval2017.zip

整理目录结构:Ultralytics 的 dataloader 要求图片放在 images/<split>/,标签放在 labels/<split>/,两者平级且同名——框架会自动将路径中的 images 替换为 labels 来查找标签:

mkdir -p images
mv val2017 images/

转换标注格式:COCO 原始标注为 JSON 格式,需转换为 YOLO 的 .txt 格式(每行:class cx cy w h,归一化坐标)。Ultralytics 自带的 convert_coco 在新版本有 bug(会误读 captions JSON),这里手动实现转换:

python -c "
import json, os
from pathlib import Path
from collections import defaultdict

with open('annotations/instances_val2017.json') as f:
    coco = json.load(f)

cat_ids = sorted([c['id'] for c in coco['categories']])
cat_map = {cid: i for i, cid in enumerate(cat_ids)}

os.makedirs('labels/val2017', exist_ok=True)

images = {img['id']: img for img in coco['images']}
anns_by_img = defaultdict(list)
for ann in coco['annotations']:
    if ann.get('iscrowd', 0):
        continue
    anns_by_img[ann['image_id']].append(ann)

for img_id, img in images.items():
    w, h = img['width'], img['height']
    txt_path = f\"labels/val2017/{Path(img['file_name']).stem}.txt\"
    with open(txt_path, 'w') as f:
        for ann in anns_by_img.get(img_id, []):
            if ann.get('bbox'):
                x, y, bw, bh = ann['bbox']
                cx, cy = (x + bw/2)/w, (y + bh/2)/h
                f.write(f\"{cat_map[ann['category_id']]} {cx:.6f} {cy:.6f} {bw/w:.6f} {bh/h:.6f}\n\")

print(f'Done: {len(images)} label files')
"

创建 coco.yaml

cat > coco.yaml << 'EOF'
path: .
val: images/val2017
names:
  0: person
  1: bicycle
  2: car
  3: motorcycle
  4: airplane
  5: bus
  6: train
  7: truck
  8: boat
  9: traffic light
  10: fire hydrant
  11: stop sign
  12: parking meter
  13: bench
  14: bird
  15: cat
  16: dog
  17: horse
  18: sheep
  19: cow
  20: elephant
  21: bear
  22: zebra
  23: giraffe
  24: backpack
  25: umbrella
  26: handbag
  27: tie
  28: suitcase
  29: frisbee
  30: skis
  31: snowboard
  32: sports ball
  33: kite
  34: baseball bat
  35: baseball glove
  36: skateboard
  37: surfboard
  38: tennis racket
  39: bottle
  40: wine glass
  41: cup
  42: fork
  43: knife
  44: spoon
  45: bowl
  46: banana
  47: apple
  48: sandwich
  49: orange
  50: broccoli
  51: carrot
  52: hot dog
  53: pizza
  54: donut
  55: cake
  56: chair
  57: couch
  58: potted plant
  59: bed
  60: dining table
  61: toilet
  62: tv
  63: laptop
  64: mouse
  65: remote
  66: keyboard
  67: cell phone
  68: microwave
  69: oven
  70: toaster
  71: sink
  72: refrigerator
  73: book
  74: clock
  75: vase
  76: scissors
  77: teddy bear
  78: hair drier
  79: toothbrush
EOF

目录结构(Ultralytics 标准布局):

/workspace/
├── images/
│   └── val2017/          # 5000张验证图片
├── labels/
│   └── val2017/          # YOLO格式标签(转换生成)
├── coco.yaml             # 数据集配置
├── yolo26s.pt            # 预训练权重
├── yolo26s.onnx          # 导出的ONNX
├── yolo26s_fp16_nhwc.mxr # 编译后的MXR
└── yolov26s_mgx_coco_8gpus.py

说明:Ultralytics 通过图片路径自动推断标签路径——将 images 替换为 labels。因此 images/val2017/000000000139.jpg 对应的标签为 labels/val2017/000000000139.txt,两者必须平级且同名。


二、Ultralytics 原生推理验证

在进入 MIGraphX 部署之前,先用 Ultralytics 自带的接口跑一遍推理,确认 DCU 环境正常、模型可用。

创建 predict.py

from ultralytics import YOLO

# 加载 YOLO26s 预训练模型(首次自动下载权重)
model = YOLO("yolo26s.pt")

# 单张图片推理
results = model.predict(
    source="https://ultralytics.com/images/bus.jpg",
    save=True,          # 保存结果图片
    device=0,           # 指定 GPU 0
)

执行:

python predict.py

看到类似输出说明 DCU 推理正常:

Ultralytics 8.4.90 🚀 Python-3.10.12 torch-2.9.0 CUDA:0 (BW1101, 147440MiB)
YOLO26s summary (fused): 122 layers, 9,496,140 parameters, 0 gradients, 20.7 GFLOPs

image 1/1 /workspace/bus.jpg: 640x480 5 persons, 1 bus, 19.4ms
Speed: 2.2ms preprocess, 19.4ms inference, 0.6ms postprocess per image at shape (1, 3, 640, 480)
Results saved to /workspace/runs/detect/predict

也可以直接用命令行:

yolo predict model=yolo26s.pt source='https://ultralytics.com/images/bus.jpg' device=0

说明:DCU 上 device=0 对应第一张卡。DTK 环境已将 HIP 映射为 CUDA 兼容接口,PyTorch 和 Ultralytics 无需额外适配即可直接使用。此步骤仅用于环境验证,后续大规模推理使用 MIGraphX 编译方案以获得更高吞吐。


三、模型导出(PT → ONNX)

创建 pt2onnx.py

from ultralytics import YOLO

# 加载 YOLO26s 预训练模型
model = YOLO("yolo26s.pt")

# 专为 MIGraphX 编译优化的导出参数
model.export(
    format="onnx",
    half=False,       # 关闭半精度导出,保持 FP32(精度转换在 MIGraphX 编译阶段完成)
    simplify=True,    # 开启常量折叠和计算图拓扑重排
    dynamic=False,    # 关闭动态 Batch Size,使用完全静态图
    batch=64,         # 写死 batch=64,与后续 MIGraphX 编译参数对齐
)

执行导出:

python pt2onnx.py

首次运行时 Ultralytics 会自动从 GitHub 下载 yolo26s.pt 权重,下载后缓存在当前目录。

导出完成后生成 yolo26s.onnx

关键点:

  1. FP32 导出 + MIGraphX 编译 FP16:刻意保持 FP32 导出 + 静态 batch,让 MIGraphX 编译器获得最优计算图结构,FP16 量化在编译阶段完成,效果更好
  2. YOLO26 NMS-Free 优势:YOLO26 默认使用一对一检测头,导出的 ONNX 计算图中不包含 NMS 节点,相比 YOLO11 减少了后处理算子,计算图更简洁,有利于 MIGraphX 的图优化和 kernel 融合

注意:如需切换回传统一对多检测头(需 NMS),可在导出时添加 end2end=False,但会引入 NMS 算子,增加部署复杂度。


四、模型编译(ONNX → MXR)

# 开启 NHWC 布局优化,性能提升明显
export MIGRAPHX_ENABLE_NHWC=1

# 编译为 FP16 离线模型
migraphx-driver compile yolo26s.onnx \
    --fp16 \
    --batch 64 \
    --gpu \
    --output yolo26s_fp16_nhwc.mxr

参数说明:

参数作用
--fp16编译为半精度,减少显存占用并提升推理速度
--batch 64指定静态 batch size,与导出阶段保持一致
--gpu指定在 GPU 上编译和运行
--output输出 .mxr 离线模型文件

编译完成后生成 yolo26s_fp16_nhwc.mxr,后续推理直接加载该文件,无需重复编译。

YOLO26 NMS-Free 对编译的影响:由于计算图中无 NMS 算子,MIGraphX 编译时无需处理动态输出形状(NMS 输出框数量不固定的问题),编译过程更稳定,推理输出 shape 固定为 (N, 300, 6),便于后续结果解析。


五、8卡并行推理测试

重要:运行推理脚本前需将 NumPy 降级到 1.25.0。前面安装 ultralytics 等依赖时可能会拉入更高版本的 NumPy,导致 MIGraphX 推理时出现兼容性问题:

pip install numpy==1.25.0

1. 推理脚本

创建 yolov26s_mgx_coco_8gpus.py

# YOLO26 🚀 by Ultralytics, adapted for AMD MIGraphX 8-GPU parallel inference
"""
Validate a trained YOLO26 model accuracy on a custom dataset with 8 GPU parallel inference
Optimized with np.memmap for zero-copy ultra-fast multiprocessing initialization.
"""

import argparse
import os
import sys
import time
import multiprocessing as mp
from pathlib import Path
import yaml
import cv2
import numpy as np
import torch
import migraphx
from tqdm import tqdm

os.environ['MIGRAPHX_ENABLE_NHWC'] = '1'

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))

from ultralytics.models.yolo.detect.val import DetectionValidator


def str2bool(v):
    if isinstance(v, bool):
        return v
    if isinstance(v, str):
        if v.lower() in ('yes', 'true', 't', 'y', '1'):
            return True
        elif v.lower() in ('no', 'false', 'f', 'n', '0'):
            return False
    raise argparse.ArgumentTypeError('Boolean value expected.')


def AllocateOutputMemory(model):
    """为模型所有输出张量预分配显存"""
    outputData = {}
    for key in model.get_outputs().keys():
        outputData[key] = migraphx.allocate_gpu(s=model.get_outputs()[key])
    return outputData


def process_images_for_gpu(all_images, start_idx, end_idx, batch_size, imgsz=640):
    """将分配给某个 GPU 的图像切分为 batch"""
    batches = []
    original_sizes = []
    num_images = end_idx - start_idx

    num_full_batches = num_images // batch_size
    remaining_images = num_images % batch_size

    for i in range(num_full_batches):
        batch_start = start_idx + i * batch_size
        batch = np.zeros((batch_size, 3, imgsz, imgsz), dtype=np.float32)
        for j in range(batch_size):
            batch[j] = all_images[batch_start + j]
        batches.append(batch)
        original_sizes.append(batch_size)

    # 不足一个 batch 的尾部数据,用全零填充补齐
    if remaining_images > 0:
        batch = np.zeros((batch_size, 3, imgsz, imgsz), dtype=np.float32)
        for j in range(remaining_images):
            batch[j] = all_images[start_idx + num_full_batches * batch_size + j]
        batches.append(batch)
        original_sizes.append(remaining_images)

    total_process = num_full_batches * batch_size + remaining_images
    return batches, total_process, original_sizes


def gpu_inference_worker(gpu_id, weights, batch_size, memmap_path, total_dataset_size,
                         start_idx, end_idx, result_queue, worker_args, imgsz=640):
    """单 GPU 推理工作进程"""
    torch.cuda.set_device(gpu_id)

    quiet_mode = worker_args.get('quiet_mode', False)
    dry_run = worker_args.get('dry_run', True)

    resultdir = os.path.join('results', f'gpu{gpu_id}')
    if not dry_run:
        os.makedirs(resultdir, exist_ok=True)

    # ---- 加载模型 ----
    if weights.split(".")[-1] == "mxr":
        try:
            model = migraphx.load(weights)
            inputName = list(model.get_inputs().keys())[0]
        except Exception as e:
            print(f"GPU {gpu_id} 加载 MXR 失败: {e}")
            return
    elif weights.split(".")[-1] == "onnx":
        max_input = {"images": [batch_size, 3, imgsz, imgsz]}
        model = migraphx.parse_onnx(weights, map_input_dims=max_input)
        inputName = model.get_parameter_names()[0]
        migraphx.quantize_fp16(model)
        model.compile(t=migraphx.get_target("gpu"), offload_copy=False, device_id=gpu_id)
    else:
        print(f"GPU {gpu_id}: 不支持的模型格式")
        return

    modelData = AllocateOutputMemory(model)

    # ---- 零拷贝读取共享内存映射 ----
    all_images = np.memmap(memmap_path, dtype=np.float32, mode='r',
                           shape=(total_dataset_size, 3, imgsz, imgsz))

    batches, total_process, original_sizes = process_images_for_gpu(
        all_images, start_idx, end_idx, batch_size, imgsz=imgsz)

    # ---- Warm up ----
    print(f"GPU {gpu_id}: Warm-up (20 iterations)...")
    warmup_batch = batches[0] if batches else np.zeros((batch_size, 3, imgsz, imgsz), dtype=np.float32)
    for _ in range(20):
        modelData[inputName] = migraphx.to_gpu(migraphx.argument(warmup_batch))
        model.run(modelData)
    print(f"GPU {gpu_id}: Warm-up completed")

    # ---- 正式推理 ----
    infer_times = []
    i = 0
    print(f"GPU {gpu_id}: Starting inference!")

    iterator = batches if quiet_mode else tqdm(batches, desc=f'GPU {gpu_id} inferencing')
    for batch_idx, batch in enumerate(iterator):
        modelData[inputName] = migraphx.to_gpu(migraphx.argument(batch))

        start = time.time()
        out = model.run(modelData)
        infer_times.append(time.time() - start)

        out = np.array(migraphx.from_gpu(out[0]))
        if not dry_run:
            original_size = original_sizes[batch_idx]
            out[:original_size].tofile(f'{resultdir}/{i}_{gpu_id}.bin')
        i += 1

    total_infer_time = sum(infer_times)
    fps = total_process / total_infer_time if total_infer_time > 0 else 0

    print(f"GPU {gpu_id}: Inference completed")
    result_queue.put({
        'gpu_id': gpu_id,
        'total_images': total_process,
        'total_infer_time': total_infer_time,
        'fps': fps
    })


def run(data, weights=None, batch_size=32, imgsz=640, conf_thres=0.001, iou_thres=0.65,
        ground_truth_json='', num_gpus=8, split_data=True, dry_run=True, quiet_mode=True):

    print(f"Starting {num_gpus}-GPU parallel inference with batch_size={batch_size}")

    # ---- 1. 解析 YAML ----
    if isinstance(data, str):
        with open(data, errors='ignore') as f:
            data_dict = yaml.safe_load(f)
    else:
        data_dict = data

    # ---- 2. 构造 YOLO26 Dataloader ----
    print("Initializing YOLO26 Dataloader...")
    validator = DetectionValidator()
    validator.args.imgsz = imgsz
    validator.args.batch_size = batch_size
    validator.args.rect = False
    validator.args.pad = 0.5
    validator.data = data_dict

    val_path = data_dict.get('val', '')
    if not os.path.isabs(val_path):
        val_path = os.path.join(data_dict.get('path', ''), val_path)

    dataloader = validator.get_dataloader(dataset_path=val_path, batch_size=batch_size)
    total_dataset_size = len(dataloader.dataset)
    print(f"Dataloader initialized! Found {total_dataset_size} images.")

    # ---- 3. 内存映射预加载 ----
    memmap_path = './yolo_perf_images.memmap'
    print(f"Preloading {total_dataset_size} images into Shared Memory Map...")

    all_images_mm = np.memmap(memmap_path, dtype=np.float32, mode='w+',
                              shape=(total_dataset_size, 3, imgsz, imgsz))

    idx = 0
    iterator = dataloader if quiet_mode else tqdm(dataloader, desc='Loading images')
    for batch in iterator:
        img = batch['img'].float() / 255.0
        img_np = img.numpy()
        num_imgs = img_np.shape[0]
        all_images_mm[idx: idx + num_imgs] = img_np
        idx += num_imgs

    all_images_mm.flush()
    print("Preload complete! Images ready for zero-copy access by 8 GPUs.")

    # ---- 4. 多进程任务分发 ----
    if split_data:
        images_per_gpu = total_dataset_size // num_gpus
    else:
        images_per_gpu = total_dataset_size

    result_queue = mp.Queue()
    processes = []

    worker_args = {
        'split_data': split_data,
        'dry_run': dry_run,
        'quiet_mode': quiet_mode,
    }

    for gpu_id in range(num_gpus):
        start_idx = gpu_id * images_per_gpu if split_data else 0
        end_idx = start_idx + images_per_gpu if split_data else total_dataset_size

        # 最后一个 GPU 处理除不尽的剩余图像
        if split_data and gpu_id == num_gpus - 1:
            end_idx = total_dataset_size

        p = mp.Process(
            target=gpu_inference_worker,
            args=(gpu_id, weights, batch_size, memmap_path, total_dataset_size,
                  start_idx, end_idx, result_queue, worker_args, imgsz)
        )
        processes.append(p)
        p.start()

    print(f"\nAll {num_gpus} processes started. Waiting for completion...\n")

    for p in processes:
        p.join()

    # ---- 5. 结果汇总 ----
    results = []
    while not result_queue.empty():
        results.append(result_queue.get())

    results.sort(key=lambda r: r['gpu_id'])

    total_fps = sum(r['fps'] for r in results)
    total_time = sum(r['total_infer_time'] for r in results) / num_gpus if num_gpus > 0 else 0

    print("\n" + "=" * 60)
    print(f"{num_gpus}-GPU INFERENCE COMPLETE")
    print("=" * 60)

    for r in results:
        print(f"GPU {r['gpu_id']}: {r['fps']:.2f} FPS ({r['total_images']} images)")
    print("-" * 60)
    print(f"Total FPS across all {num_gpus} GPUs: {total_fps:.2f} samples/s")
    print(f"Average inference time per GPU: {total_time:.4f}s")
    print("=" * 60)

    # 清理临时内存映射文件
    if os.path.exists(memmap_path):
        try:
            os.remove(memmap_path)
            print("Cleaned up temporary memory map file.")
        except Exception as e:
            print(f"Warning: Failed to clean up {memmap_path}: {e}")


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', type=str, default='./coco.yaml', help='dataset.yaml path')
    parser.add_argument('--weights', type=str, default='./yolo26s_fp16_nhwc.mxr', help='model.mxr or model.onnx path')
    parser.add_argument('--batch-size', type=int, default=64, help='batch size')
    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--num-gpus', type=int, default=8, help='number of GPUs to use')
    parser.add_argument('--split-data', type=str2bool, default=True, help='True=split dataset evenly across GPUs')
    parser.add_argument('--dry-run', type=str2bool, default=True, help='True=do not save result files')
    parser.add_argument('--quiet-mode', type=str2bool, default=True, help='True=minimal output')
    parser.add_argument('--ground_truth_json', type=str, default='', help='annotation file path')
    opt = parser.parse_args()
    return opt


def main(opt):
    run(**vars(opt))


if __name__ == "__main__":
    mp.set_start_method('spawn', force=True)
    opt = parse_opt()
    main(opt)

脚本核心设计:

模块说明
np.memmap 预加载将全部图像一次性写入共享内存文件,8个子进程零拷贝读取,避免串行加载卡顿
mp.Process 多进程每个 GPU 一个独立进程,spawn 启动方式保证 ROCm 兼容性
Warm-up 20轮预热消除首次 kernel 编译开销,保证测时准确
数据均分8卡按 total // 8 均分,最后一张卡兜底处理余数

2. 运行脚本

创建 run_yolov26s_mgx_coco.sh

#!/bin/bash

# MIGraphX 性能优化环境变量
export MIGRAPHX_ENABLE_NHWC=1        # 开启 NHWC 布局,性能提升明显
export MIGRAPHX_ENABLE_CUTLASS=1    # 开启 CUTLASS 算子
export HIP_ALLOC_INITIALIZE=0        # 跳过显存初始化,减少分配开销
export HIP_USE_GRAPH_QUEUE_POOL=1   # 使用 Graph 队列池
export HIP_KERNEL_EVENT_SYSTENFENCE=1
export HAS_FORCE_FINE_GRAIN_PCIE=1  # 强制细粒度 PCIe 传输

log_dir=logs
mkdir -p ${log_dir}

python yolov26s_mgx_coco_8gpus.py \
    --data ./coco.yaml \
    --weights ./yolo26s_fp16_nhwc.mxr \
    --batch-size 64 \
    2>&1 | tee "${log_dir}/yolo_fp16_$(date +%Y-%m-%d_%H-%M-%S).log"

3. 执行测试

chmod +x run_yolov26s_mgx_coco.sh
bash run_yolov26s_mgx_coco.sh

六、测试结果

============================================================
8-GPU INFERENCE COMPLETE
============================================================
GPU 0: 1650.78 FPS (625 images)
GPU 1: 1651.98 FPS (625 images)
GPU 2: 1649.57 FPS (625 images)
GPU 3: 1647.60 FPS (625 images)
GPU 4: 1651.12 FPS (625 images)
GPU 5: 1647.96 FPS (625 images)
GPU 6: 1644.72 FPS (625 images)
GPU 7: 1647.57 FPS (625 images)
------------------------------------------------------------
Total FPS across all 8 GPUs: 13191.31 samples/s
Average inference time per GPU: 0.3790s
============================================================
指标数值
单卡平均 FPS≈ 1650 FPS
8卡总吞吐≈ 13191 samples/s
单卡平均推理耗时≈ 0.379s(625张/卡)
单张推理耗时(满 batch)≈ 0.61 ms
Batch Size64
精度FP16
数据集COCO Val2017(5000张,每卡625张)

性能对比:

推理方式单张耗时单卡 FPS加速比
Ultralytics 原生(PT, FP32)19.4 ms~51
MIGraphX FP16(batch=64)0.61 ms~165032×

分析:从原生 Ultralytics 推理的 19.4ms/张降至 MIGraphX FP16 编译后的 0.61ms/张,加速达 32倍。加速来源:① FP16 量化减少显存带宽和计算量;② MIGraphX 编译期图优化(算子融合、常量折叠);③ batch=64 充分利用 DCU 并行度;④ YOLO26 NMS-Free 架构消除了后处理同步开销。8卡总吞吐 13191 samples/s,扩展效率接近线性。


七、总结

本文完整记录了在海光 BW1100 DCU 上使用 MIGraphX 进行 YOLO26s 8卡并行推理的全流程,核心要点:

  1. YOLO26 NMS-Free 部署优势:默认一对一检测头导出的 ONNX 计算图中无 NMS 算子,输出 shape 固定为 (N, 300, 6),MIGraphX 编译更稳定,无需处理动态输出形状
  2. 模型导出策略:PT 导出时保持 FP32 + 静态 batch,FP16 量化交给 MIGraphX 编译阶段完成,计算图结构最优
  3. MIGraphX 编译优化--fp16 + MIGRAPHX_ENABLE_NHWC=1 是性能关键组合,NHWC 布局可显著提升卷积算子效率
  4. 多进程零拷贝:通过 np.memmap 将全部图像预加载至共享内存,8个子进程零拷贝读取,消除了数据加载瓶颈
  5. 8卡并行扩展:单卡 ~1650 FPS,8卡总吞吐 ~13191 samples/s,单张推理仅 0.61ms,扩展效率接近线性

适用场景:该方案适用于大规模图像检测的批量推理任务,如工业质检、安防监控等需要高吞吐的部署环境。


相关链接

作者注:本文涉及的内容仍处于持续学习与验证阶段,个人理解仅供参考。如有疏漏或更好的实践经验,欢迎大家补充讨论。

Logo

北京人形旗下天工造物具身智能开源社区,聚焦具身天工与慧思开物两大平台

更多推荐