【AMD ROCm 实战】云端 AI 开发系列（二）：CUDA 到 ROCm 迁移实战——YOLOv8 目标检测模型全流程对标

# 安装 YOLOv8
pip install ultralytics opencv-python
# 验证安装
python -c "from ultralytics import YOLO; print('✅ YOLOv8 installed successfully!')"

# download_model.py
from ultralytics import YOLO
# 下载 YOLOv8n (nano 版本，适合边缘部署)
model_n = YOLO('yolov8n.pt')
print(f"✅ YOLOv8n downloaded: {model_n.info()}")
# 下载 YOLOv8m (medium 版本，平衡精度与速度)
model_m = YOLO('yolov8m.pt')
print(f"✅ YOLOv8m downloaded: {model_m.info()}")
# 下载 YOLOv8l (large 版本，高精度)
model_l = YOLO('yolov8l.pt')
print(f"✅ YOLOv8l downloaded: {model_l.info()}")

# yolo_cuda_inference.py
from ultralytics import YOLO
import torch
import time
# 加载模型
model = YOLO('yolov8m.pt')
# 指定 GPU (NVIDIA)
device = 'cuda:0'
model.to(device)
print(f"Using device: {torch.cuda.get_device_name(0)}")
# 准备测试图片
image_path = 'test_image.jpg'
# 推理
results = model.predict(image_path, device=device, conf=0.25)
# 输出结果
for r in results:
    print(f"Detected {len(r.boxes)} objects")

# yolo_rocm_inference.py
from ultralytics import YOLO
import torch
import time
print("=" * 60)
print("YOLOv8 Inference on AMD ROCm")
print("=" * 60)
# 1. 检查 ROCm 是否可用
print(f"\n✅ ROCm available: {torch.cuda.is_available()}")
if not torch.cuda.is_available():
    raise RuntimeError("ROCm not available! Please check your installation.")
# 2. 加载模型
model = YOLO('yolov8m.pt')
# 3. 指定 GPU (AMD MI300X)
# 💡 关键：ROCm 依然使用 'cuda' 接口，无需修改代码！
device = 'cuda:0'
model.to(device)
print(f"🎯 Using device: {torch.cuda.get_device_name(0)}")
print(f"💾 Total memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
# 4. 准备测试图片
image_path = 'test_image.jpg'
# 5. 预热 (Warm-up)
print("\n🔥 Warming up...")
for _ in range(10):
    model.predict(image_path, device=device, verbose=False)
# 6. 正式测试 (100 次推理)
num_iterations = 100
print(f"\n🚀 Running {num_iterations} inference iterations...")
start_time = time.time()
for i in range(num_iterations):
    results = model.predict(image_path, device=device, verbose=False, conf=0.25)
end_time = time.time()
total_time = end_time - start_time
avg_latency = (total_time / num_iterations) * 1000
fps = num_iterations / total_time
print(f"\n📈 Results:")
print(f"   Total time: {total_time:.2f} seconds")
print(f"   Average latency: {avg_latency:.2f} ms")
print(f"   FPS: {fps:.2f}")
print(f"   Detected objects per frame: {len(results[0].boxes)}")
print("\n" + "=" * 60)

RuntimeError: HIP error: invalid device ordinal

# 1. 检查可用的 GPU
rocm-smi --showproductname
# 2. 设置正确的 GPU 编号（从 0 开始）
export HIP_VISIBLE_DEVICES=0
# 3. 在 Python 中验证
python -c "import torch; print(torch.cuda.device_count())"
# 应输出: 1

# 方案 1: 调整 NMS 阈值
results = model.predict(
    image_path, 
    device='cuda:0', 
    conf=0.25,
    iou=0.7  # 提高 IoU 阈值，减少重复框
)
# 方案 2: 手动后处理
import torchvision
def custom_nms(boxes, scores, iou_threshold=0.7):
    """自定义 NMS，确保跨平台一致性"""
    keep_indices = torchvision.ops.nms(boxes, scores, iou_threshold)
    return keep_indices
# 在推理后应用
for r in results:
    boxes = r.boxes.xyxy
    scores = r.boxes.conf
    keep = custom_nms(boxes, scores, iou_threshold=0.7)
    r.boxes = r.boxes[keep]

# 方案 1: 固定输入尺寸
# 导出模型时指定固定输入尺寸
from ultralytics import YOLO
model = YOLO('yolov8m.pt')
model.export(format='onnx', imgsz=640, dynamic=False)  # 禁用动态形状
# 方案 2: 定期清理缓存
class YOLOROCMInference:
    def __init__(self, model_path):
        self.model = YOLO(model_path)
        self.model.to('cuda:0')
        self.inference_count = 0
    def predict(self, image_path, **kwargs):
        results = self.model.predict(image_path, device='cuda:0', **kwargs)
        self.inference_count += 1
        # 每 500 次推理清理一次缓存
        if self.inference_count % 500 == 0:
            torch.cuda.empty_cache()
            print(f"♻️  Cache cleared at inference #{self.inference_count}")
        return results
# 使用示例
inference_engine = YOLOROCMInference('yolov8m.pt')
for i in range(1000):
    results = inference_engine.predict(f'image_{i}.jpg')

# export_to_onnx.py
from ultralytics import YOLO
# 加载模型
model = YOLO('yolov8m.pt')
# 导出为 ONNX
model.export(
    format='onnx',
    imgsz=640,
    dynamic=False,  # 固定输入尺寸，提升性能
    simplify=True   # 简化计算图
)
print("✅ ONNX model exported: yolov8m.onnx")

# 安装 ONNX Runtime ROCm
pip uninstall onnxruntime
pip install onnxruntime-rocm

# onnx_rocm_inference.py
import onnxruntime as ort
import numpy as np
import cv2
import time
print("=" * 60)
print("YOLOv8 ONNX Inference on AMD ROCm")
print("=" * 60)
# 1. 创建 Session
session_options = ort.SessionOptions()
session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
session = ort.InferenceSession(
    'yolov8m.onnx',
    session_options,
    providers=['ROCMExecutionProvider', 'CPUExecutionProvider']
)
print(f"🎯 Using provider: {session.get_providers()[0]}")
# 2. 预处理函数
def preprocess_image(image_path, input_size=640):
    img = cv2.imread(image_path)
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img_resized = cv2.resize(img_rgb, (input_size, input_size))
    img_normalized = img_resized.astype(np.float32) / 255.0
    img_transposed = np.transpose(img_normalized, (2, 0, 1))
    img_batch = np.expand_dims(img_transposed, axis=0)
    return img_batch
# 3. 推理测试
image_path = 'test_image.jpg'
input_data = preprocess_image(image_path)
# 预热
for _ in range(10):
    session.run(None, {'images': input_data})
# 正式测试
num_iterations = 100
start_time = time.time()
for _ in range(num_iterations):
    output = session.run(None, {'images': input_data})
end_time = time.time()
avg_latency = (end_time - start_time) / num_iterations * 1000
fps = num_iterations / (end_time - start_time)
print(f"\n📈 Results:")
print(f"   Average latency: {avg_latency:.2f} ms")
print(f"   FPS: {fps:.2f}")
print("\n" + "=" * 60)