Skip to content

Commit 28aede1

Browse files
committed
fix(rocm): use PyTorch+HIP for inference instead of ONNX
Ultralytics' ONNX loader only supports CUDAExecutionProvider (NVIDIA). On ROCm, it falls back to CPU even though ROCMExecutionProvider is available. PyTorch + HIP runs natively on AMD GPUs via device='cuda'. - Change ROCm BackendSpec: onnx → pytorch (skip ONNX export entirely) - Set YOLO_AUTOINSTALL=0 in detect.py to prevent ultralytics from auto-installing onnxruntime-gpu (NVIDIA) at runtime
1 parent 2d32d52 commit 28aede1

File tree

3 files changed

+10
-6
lines changed

3 files changed

+10
-6
lines changed

skills/detection/yolo-detection-2026/scripts/detect.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,16 @@
1515
"""
1616

1717
import sys
18+
import os
1819
import json
1920
import argparse
2021
import signal
2122
import time
2223
from pathlib import Path
2324

25+
# Prevent ultralytics from auto-installing packages (e.g. onnxruntime-gpu on ROCm)
26+
os.environ.setdefault("YOLO_AUTOINSTALL", "0")
27+
2428
# Import env_config — try multiple locations:
2529
# 1. Same directory as detect.py (bundled copy)
2630
# 2. DeepCamera repo: skills/lib/

skills/detection/yolo-detection-2026/scripts/env_config.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -51,9 +51,9 @@ class BackendSpec:
5151
),
5252
"rocm": BackendSpec(
5353
name="rocm",
54-
export_format="onnx",
55-
model_suffix=".onnx",
56-
half=False, # ONNX Runtime ROCm handles precision internally
54+
export_format="pytorch", # PyTorch + HIP — ultralytics ONNX doesn't support ROCMExecutionProvider
55+
model_suffix=".pt",
56+
half=False,
5757
),
5858
"mps": BackendSpec(
5959
name="mps",

skills/lib/env_config.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -51,9 +51,9 @@ class BackendSpec:
5151
),
5252
"rocm": BackendSpec(
5353
name="rocm",
54-
export_format="onnx",
55-
model_suffix=".onnx",
56-
half=False, # ONNX Runtime ROCm handles precision internally
54+
export_format="pytorch", # PyTorch + HIP — ultralytics ONNX doesn't support ROCMExecutionProvider
55+
model_suffix=".pt",
56+
half=False,
5757
),
5858
"mps": BackendSpec(
5959
name="mps",

0 commit comments

Comments
 (0)