【Yolo v8 window系统使用AMD GPU进行训练】

最新推荐文章于 2024-08-24 22:36:08 发布

买榴莲也用卷

最新推荐文章于 2024-08-24 22:36:08 发布

阅读量89

点赞数

文章标签： YOLO python 开发语言

本文链接：https://blog.csdn.net/GG9527li/article/details/141326165

版权

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