使用 Ultralytics YOLO 和 ROCm 训练 AMD GPU#

Ultralytics 支持通过 PyTorch ROCm 在兼容的 AMD GPU 上进行训练、验证和推理。PyTorch 有意通过与 CUDA 相同的 torch.cuda Python API 暴露 ROCm 设备，因此 Ultralytics 不需要为原生 PyTorch 模型单独使用 rocm 设备类型。安装 ROCm 版本的 PyTorch，然后使用标准的 device=0 或 device=cuda:0 语法选择 AMD GPU。

AMD 还提供了独立于 PyTorch ROCm 的推理技术。支持某一 AMD 产品并不意味着 支持所有 AMD 运行时或加速器。

支持概览#

此表介绍用于训练、验证、导出和预测的 Ultralytics Python 包。

ROCm 为何使用 CUDA 设备名称#

PyTorch 的 ROCm 构建在内部使用 HIP，但会有意复用 torch.cuda 接口。例如， torch.cuda.is_available()、torch.cuda.device_count() 和 torch.cuda.get_device_name() 可与受支持的 AMD GPU 配合使用。 这种设计使同一套 Ultralytics 训练路径能够服务于 NVIDIA CUDA 和 AMD ROCm，而无需重复的后端。

详情请参阅官方的 PyTorch HIP 语义。

安装 PyTorch ROCm#

1. 确认你的操作系统和 GPU 出现在 ROCm 兼容性矩阵中。

2. 使用 PyTorch 安装选择器 选择与 已安装 ROCm 版本匹配的 ROCm 构建。

3. 在 PyTorch 之后安装 Ultralytics：

   pip install ultralytics

4. 验证 PyTorch 是否能识别 AMD GPU：

   import torch
   
   print(torch.cuda.is_available())
   print(torch.cuda.get_device_name(0))
   print(torch.version.hip)

torch.cuda.is_available() 应返回 True，设备名称应标识出你的 AMD GPU，并且 torch.version.hip 应报告 ROCm 构建提供的 HIP 版本。

在 AMD GPU 上训练#

使用与标准 Ultralytics Train 模式 相同的设备参数。

from ultralytics import YOLO

model = YOLO("yolo26n.pt")

# Train on the first AMD GPU exposed by PyTorch ROCm
results = model.train(data="coco8.yaml", epochs=100, imgsz=640, device=0)

# Train across two AMD GPUs
results = model.train(data="coco8.yaml", epochs=100, imgsz=640, device=[0, 1])

验证 和 预测 使用相同的设备选择：

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

AMP 兼容性#

Ultralytics 默认启用 AMP，并在训练前比较全精度和混合精度结果。如果检查 检测到结果不兼容，则会禁用 AMP，以防止出现 NaN 损失或零 mAP 训练。ROCm AMP 的行为可能会随 PyTorch 和 ROCm 版本变化，因此排查特定技术栈故障时请使用 amp=False：

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

尚不支持的功能#

MIGraphX#

MIGraphX 是 AMD 的图优化和推理运行时。当前 Ultralytics Python 包不包含原生 MIGraphX 模型 加载功能。活跃的实现、AMD 容器、依赖项、测试和文档都在 PR #24137 中统一跟踪。在该工作合并并通过 AMD 硬件 验证之前，不应将导出 ONNX 模型描述为 Python 包中的原生 MIGraphX 支持。

DirectML#

Ultralytics Python 包没有用于 Windows 训练或预测的 DirectML 后端。DirectML 不同于 ROCm，正常运行的 ROCm 环境不会启用 device=directml。

Ryzen AI NPU#

Ryzen AI NPU 不会通过 PyTorch ROCm 暴露，也不是原生的 Ultralytics 设备。社区工作流可能会 将 YOLO 模型导出为 ONNX，并使用 AMD 的外部 Ryzen AI 或 Vitis AI 工具运行，但该运行时、转换过程和 硬件兼容性不属于 Ultralytics 支持的执行路径。

总结#

对于受支持的 AMD GPU 训练、验证和推理，请使用兼容的 PyTorch ROCm 构建以及 device=0 或 device=cuda:0。 将 MIGraphX、DirectML 和 Ryzen AI NPU 视为独立功能：仅安装 ROCm 或导出 ONNX 模型不会启用其中任何一项。

常见问题#