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
Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐