使用 Ultralytics YOLO、ROCm 和 MIGraphX 在 AMD GPU 上训练和推理#

Ultralytics YOLO 可通过 AMD 的开放 GPU 计算栈 ROCm 在 AMD Instinct 和受支持的 Radeon GPU 上运行。使用 ROCm 版 PyTorch 后，你可以像在其他任何 GPU 上一样，使用相同的 device=0 参数对 .pt 模型进行训练、验证和预测，无需更改代码。

部署时，将训练好的模型导出为 ONNX，并通过 AMD 的计算图优化和推理引擎 MIGraphX 运行。Ultralytics ONNX 后端会检测 ROCm，并自动选择 ONNX Runtime MIGraphXExecutionProvider，因此在 NVIDIA GPU 上运行的同一个 predict 调用也能在 AMD GPU 上运行。Ultralytics 定期运行的持续集成会在 AMD GPU 硬件上针对每个任务执行 ONNX 导出和 MIGraphX 推理。

安装#

Python 技术栈完全通过 pip 从 AMD 的 ROCm wheel 索引安装，因此无需 apt 软件包或 root 权限即可安装 PyTorch、插件或 ROCm 运行时库，前提是主机已安装 AMD GPU 内核驱动程序（amdgpu / /dev/kfd）。

# Read the GPU architecture from the amdgpu driver, e.g. device-gfx1151
GPU_ARCH=$(grep -rhs '^gfx_target_version' /sys/class/kfd/kfd/topology/nodes | awk '$2 {printf "device-gfx%d%d%x\n", int($2/10000), int($2/100)%100, $2%100}' | sort -u | paste -sd, -)

# ROCm PyTorch + torchvision for the detected architecture
pip install "torch[${GPU_ARCH:?No AMD GPU found}]" "torchvision[$GPU_ARCH]" --index-url https://stable.repo.amd.com/rocm/whl-next/

# Ultralytics
pip install ultralytics

Ultralytics 会在首次导出 ONNX 时自动安装 onnx。在 ROCm 系统上首次进行 ONNX 推理时，它还会安装 MIGraphX 执行提供程序插件（onnxruntime-ep-migraphx）及其加载的 MIGraphX 运行时库（migraphx-libs），并固定为经过组合测试的版本。该插件仅面向 Linux x86_64 上的 Python 3.11 或更高版本发布；在其他 Python 版本或架构上，或者使用 YOLO_AUTOINSTALL=false 且未安装插件时，ONNX 推理会发出警告并回退到 CPU。详细说明和最佳实践请参阅 YOLO26 安装指南；如果遇到任何问题，请查阅常见问题指南。

训练或导出前，请确认 PyTorch 能检测到 AMD GPU：

import torch

print(torch.cuda.is_available())  # 正常运行的 ROCm 安装会返回 True
print(torch.cuda.get_device_name(0))  # AMD GPU 名称，例如 AMD Radeon 8060S
print(torch.version.hip)  # ROCm PyTorch 版本的 HIP 版本；CUDA 或 CPU 版本则为 None

使用 PyTorch ROCm 在 AMD GPU 上训练#

基于 .pt 模型的原生训练、验证和预测通过 PyTorch ROCm 在 AMD GPU 上运行，与下文介绍的 MIGraphX 推理路径无关。按照安装指南安装 ROCm 版 PyTorch，然后像运行其他任何 Ultralytics 训练任务一样使用相同的设备参数：

from ultralytics import YOLO

model = YOLO("yolo26n.pt")

# 使用 PyTorch ROCm 检测到的第一块 AMD GPU 进行训练
results = model.train(data="coco8.yaml", epochs=100, imgsz=640, device=0)

# 使用两块 AMD GPU 进行训练
results = model.train(data="coco8.yaml", epochs=100, imgsz=640, device=[0, 1])

.pt 模型的验证和预测使用相同的设备选择方式：

yolo detect val model=path/to/best.pt data=coco8.yaml device=0
yolo predict model=path/to/best.pt source=path/to/image.jpg device=0

yolo detect train data=coco8.yaml model=yolo26n.pt device=0 amp=False

使用 MIGraphX 导出和推理#

ONNX Runtime 是一个跨平台推理引擎，可通过可插拔的执行提供程序 (EP) 在多种硬件后端上运行同一个 ONNX 模型。每个 EP 都会将 ONNX 计算图映射到特定加速器上：NVIDIA GPU 使用 CUDA，Apple silicon 使用 CoreML，而 ROCm 上的 AMD GPU 使用 MIGraphX。MIGraphX 以 onnxruntime-ep-migraphx 插件形式提供，并会将计算图编译为针对你的 GPU 调优的程序。

MIGraphX 推理使用标准的 ONNX 导出。将模型导出为 ONNX（format="onnx"）后，Ultralytics 推理会自动使用 MIGraphXExecutionProvider 后端，在 AMD GPU 上运行 ONNX 模型。

MIGraphX 推理的主要功能#

• 自动选择提供程序：在 ROCm 主机上，ONNX 后端会注册插件并选择 MIGraphXExecutionProvider，无需更改代码。如果插件缺失或加载失败，推理会发出警告并回退到 CPU。
• 计算图优化：MIGraphX 会执行算子融合、内存规划和内核选择，并针对 AMD GPU 架构进行调优。
• 零拷贝 IO 绑定：通过 DLPack 协议将输入和输出直接绑定到 GPU 张量，避免推理期间与主机之间的数据往返传输。
• 精度选项：使用 FP32 运行，或导出 FP16 ONNX 模型以进行低精度推理。
• 可移植的产物：单个 .onnx 文件即可在 CPU、NVIDIA GPU 和 AMD GPU 上运行，让你通过一次导出面向多个平台。
• 可复现的部署：完整技术栈（ROCm PyTorch、MIGraphX 插件及其库）均通过 pip 从 AMD 的 ROCm wheel 索引安装。

支持的任务#

MIGraphX 推理支持 Ultralytics 的全部七项任务。语义分割和深度估计仅适用于 YOLO26，这是唯一提供这些任务头的系列。

用法#

开始阅读使用说明前，请务必了解 Ultralytics 提供的 YOLO26 模型，以便根据项目需求选择最合适的模型。

ONNX 格式支持导出、预测和验证模式。要在 AMD GPU 上进行推理和验证，需要 ROCm 系统并安装 MIGraphX 插件。导出模型后，加载导出的模型，即可在 device=0 上运行推理或验证模型精度。

from ultralytics import YOLO

# 加载 YOLO26 模型
model = YOLO("yolo26n.pt")

# 将模型导出为 ONNX 格式
model.export(format="onnx")  # 创建 'yolo26n.onnx'

from ultralytics import YOLO

# 加载导出的 ONNX 模型，并在 AMD GPU 上运行推理
model = YOLO("yolo26n.onnx")

# 在 ROCm 上会自动选择 MIGraphX 执行提供程序
results = model.predict("https://ultralytics.com/images/bus.jpg", device=0)

from ultralytics import YOLO

# 加载导出的 ONNX 模型
model = YOLO("yolo26n.onnx")

# 在 AMD GPU 上使用 COCO8 数据集验证准确率
metrics = model.val(data="coco8.yaml", device=0)

导出参数#

MIGraphX 推理会复用 ONNX 导出参数。AMD GPU 部署最相关的选项包括：

完整的导出参数列表请参阅 ONNX 集成和 Ultralytics 导出文档页面。

使用 MIGraphX 在 AMD GPU 上部署#

如需开箱即用的环境，Dockerfile-amd 会构建预装 ROCm PyTorch 和 MIGraphX EP 的 ultralytics/ultralytics:latest-amd 镜像。请参阅 Docker 快速入门中的使用 GPU，了解用于向容器开放 AMD GPU 的 docker run 标志。

基准测试#

以下 YOLO26 基准测试由 Ultralytics 团队在 AMD Radeon 8060S GPU (Ryzen AI Max+ PRO 395) 上运行，对比了 PyTorch 与 ONNX 在 MIGraphX 执行提供程序上的速度和准确率。两种格式都以 FP32 精度在 GPU 上运行：PyTorch 使用 ROCm，ONNX 使用 MIGraphX。

支持情况一览#

支持某一种 AMD 产品或运行时，并不代表支持所有 AMD 加速器。此表汇总了 Ultralytics Python 包当前的支持状态。

总结#

在本指南中，你学习了如何使用标准的 device=0 参数，通过 PyTorch ROCm 在 AMD GPU 上训练 Ultralytics YOLO26 模型，然后将模型导出为 ONNX，以便通过 ONNX Runtime MIGraphX 执行提供程序加速推理。ONNX 后端会自动选择 MIGraphXExecutionProvider，缓存编译后的程序以加快后续加载，并且无需修改代码即可支持所有 YOLO26 任务。

如需了解其他部署目标，请浏览集成指南页面，并通过基准测试模式比较导出格式。

常见问题#