使用 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 训练模式相同的设备参数。

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])

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

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 模型都不会启用其中任何一项。

常见问题#