pip install onnxruntime_directml
将你的YOLO model转换成 ONNX格式
from ultralytics import YOLO
model = YOLO('yolov8n.pt')
model.export(format='onnx')
添加 ‘DmlExecutionProvider’ 这个字符串到指定的包对应的行中 133 to 140 in “venv\Lib\site-packages\ultralytics\nn\autobackend.py”:
elif onnx: # ONNX Runtime
LOGGER.info(f'Loading {w} for ONNX Runtime inference...')
# check_requirements(('onnx', 'onnxruntime-gpu' if cuda else 'onnxruntime'))
import onnxruntime
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['DmlExecutionProvider', 'CPUExecutionProvider']
session = onnxruntime.InferenceSession(w, providers=providers)
output_names = [x.name for x in session.get_outputs()]
metadata = session.get_modelmeta().custom_metadata_map # metadata